Fix bugs

main.py

@@ -168,7 +168,9 @@ def main_worker(gpu, ngpus_per_node, args):
 
     pre_train_dataset = pcl.loader.PreImager(pre_train_dir, eval_augmentation)
 
-    train_dataset = pcl.loader.ImageFolderInstance(train_dir, eval_augmentation)
+    train_dataset = pcl.loader.ImageFolderInstance(
+        train_dir,
+        pcl.loader.TwoCropsTransform(eval_augmentation))
     eval_dataset = pcl.loader.ImageFolderInstance(
         train_dir,
         eval_augmentation)
@@ -195,7 +197,7 @@ def main_worker(gpu, ngpus_per_node, args):
         drop_last=True)
 
     train_loader = torch.utils.data.DataLoader(
-        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
+        train_dataset, batch_size=args.batch_size//2, shuffle=(train_sampler is None),
         num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True)
 
     # dataloader for center-cropped images, use larger batch size to increase speed
@@ -275,15 +277,18 @@ def train(train_loader, model, criterion, optimizer, epoch, args, cluster_result
 
         im_q = []
         im_k = []
+        if cluster_result is None:
             class_number = len(images)
             class_len = len(images[0])
             for _i in range(0, class_len, 2):
                 for c in range(class_number):
                     im_q.append(images[c][_i])
                     im_k.append(images[c][_i+1])
-
             im_q = torch.stack(im_q)
             im_k = torch.stack(im_k)
+        else:
+            im_q = images[0]
+            im_k = images[1]
 
         if args.gpu is not None:
             im_q = im_q.cuda(args.gpu, non_blocking=True)

pcl/loader.py

@@ -22,11 +22,12 @@ class PreImager(tud.Dataset):
 
         self.aug = aug
         self.images = img_class
+        self.loader = data_meta.loader
 
     def __getitem__(self, index):
         imgs = []
         for i in range(self.class_number):
-            img = image.read_image(self.images[i][index]).float()
+            img = self.loader(self.images[i][index])
             out = self.aug(img)
             imgs.append(out)
         return imgs, index