Initial commit

IdentitySubset.py

@@ -0,0 +1,13 @@
+
+import torch
+
+class IdentitySubset(torch.utils.data.Dataset):
+    def __init__(self, full_ds, indices, id_mapping):
+        self.full_ds = full_ds
+        self.indices = indices
+        self.id_mapping = id_mapping
+    def __getitem__(self, idx):
+        img, old_id = self.full_ds[self.indices[idx]]
+        return img, self.id_mapping[old_id.item()]
+    def __len__(self):
+        return len(self.indices)