Simclr (#33)

* remove jitter * remove jitter * remove jitter * remove jitter * remove jitter * remove jitter * remove jitter * remove jitter * added hparams test * added hparams test * added hparams test * added hparams test * added hparams test * added hparams test * training batch clean up
Lightning-Universe · Jun 8, 2020 · 977f2ed · 977f2ed
1 parent b7ab0ef
commit 977f2ed
Show file tree

Hide file tree

Showing 11 changed files with 910 additions and 99 deletions.
diff --git a/pl_bolts/datamodules/imagenet_dataset.py b/pl_bolts/datamodules/imagenet_dataset.py
@@ -20,6 +20,7 @@ def __init__(self, root, split='train',
                  num_classes=-1,
                  num_imgs_per_class=-1,
                  num_imgs_per_class_val_split=50,
+                 meta_root=None,
                  **kwargs):
         """
         Official train set gets split into train, val. (using nb_imgs_per_val_class for each class).
@@ -47,7 +48,8 @@ def __init__(self, root, split='train',
 
         self.split = split
         split_root = os.path.join(root, split)
-        wnid_to_classes = load_meta_file(split_root)[0]
+        meta_root = meta_root if meta_root is not None else split_root
+        wnid_to_classes = load_meta_file(meta_root)[0]
 
         super(ImageNet, self).__init__(self.split_folder, **kwargs)
         self.root = root

diff --git a/pl_bolts/datamodules/ssl_imagenet_dataloaders.py b/pl_bolts/datamodules/ssl_imagenet_dataloaders.py
@@ -12,11 +12,13 @@ class SSLImagenetDataLoaders(BoltDataLoaders):
 
     def __init__(self,
                  data_dir,
+                 meta_root=None,
                  num_workers=16):
 
         super().__init__()
         self.data_dir = data_dir
         self.num_workers = num_workers
+        self.meta_root = meta_root
 
     @property
     def num_classes(self):
@@ -58,6 +60,7 @@ def train_dataloader(self, batch_size, num_images_per_class=-1, transforms=None,
 
         dataset = UnlabeledImagenet(self.data_dir,
                                     num_imgs_per_class=num_images_per_class,
+                                    meta_root=self.meta_root,
                                     split='train',
                                     transform=transforms)
         loader = DataLoader(
@@ -76,6 +79,7 @@ def val_dataloader(self, batch_size, num_images_per_class=50, transforms=None, a
 
         dataset = UnlabeledImagenet(self.data_dir,
                                     num_imgs_per_class_val_split=num_images_per_class,
+                                    meta_root=self.meta_root,
                                     split='val',
                                     transform=transforms)
         loader = DataLoader(
@@ -93,6 +97,7 @@ def test_dataloader(self, batch_size, num_images_per_class, transforms=None, add
 
         dataset = UnlabeledImagenet(self.data_dir,
                                     num_imgs_per_class=num_images_per_class,
+                                    meta_root=self.meta_root,
                                     split='test',
                                     transform=transforms)
         loader = DataLoader(

diff --git a/pl_bolts/datamodules/submit.py b/pl_bolts/datamodules/submit.py