add tests for USPS dataset (#3466)

Summary: Co-authored-by: Vasilis Vryniotis <[email protected]>

Reviewed By: fmassa

Differential Revision: D26756275

fbshipit-source-id: 52ed0b84541cde380a3ac539cc49c92d57bbebab

Author: zhangguanheng66

test/test_datasets.py

@@ -23,6 +23,7 @@
 import shutil
 import json
 import random
+import bz2
 import torch.nn.functional as F
 import string
 import io
@@ -1173,5 +1174,24 @@ def inject_fake_data(self, tmpdir, config):
         return num_images
 
 
+class USPSTestCase(datasets_utils.ImageDatasetTestCase):
+    DATASET_CLASS = datasets.USPS
+
+    CONFIGS = datasets_utils.combinations_grid(train=(True, False))
+
+    def inject_fake_data(self, tmpdir, config):
+        num_images = 2 if config["train"] else 1
+
+        images = torch.rand(num_images, 256) * 2 - 1
+        labels = torch.randint(1, 11, size=(num_images,))
+
+        with bz2.open(pathlib.Path(tmpdir) / f"usps{'.t' if not config['train'] else ''}.bz2", "w") as fh:
+            for image, label in zip(images, labels):
+                line = " ".join((str(label.item()), *[f"{idx}:{pixel:.6f}" for idx, pixel in enumerate(image, 1)]))
+                fh.write(f"{line}\n".encode())
+
+        return num_images
+
+
 if __name__ == "__main__":
     unittest.main()