Update code

Author: lizekai-richard

dd_ranking/metrics/general.py

@@ -94,6 +94,7 @@ def __init__(self,
         channel, im_size, num_classes, dst_train, dst_test, class_map, class_map_inv = get_dataset(dataset, 
                                                                                                    real_data_path, 
                                                                                                    im_size,
+                                                                                                   custom_train_trans,
                                                                                                    custom_val_trans,
                                                                                                    use_zca)
         self.num_classes = num_classes

dd_ranking/metrics/hard_label.py

@@ -54,6 +54,7 @@ def __init__(self, config: Config=None, dataset: str='CIFAR10', real_data_path:
                                                                                                    real_data_path, 
                                                                                                    im_size, 
                                                                                                    use_zca,
+                                                                                                   custom_train_trans,
                                                                                                    custom_val_trans,
                                                                                                    device)
         self.images_train, self.labels_train, self.class_indices_train = self.load_real_data(dst_train, class_map, num_classes)

dd_ranking/metrics/soft_label.py

@@ -23,7 +23,7 @@ def __init__(self, config: Config=None, dataset: str='CIFAR10', real_data_path:
                  soft_label_criterion: str='kl', data_aug_func: str='cutmix', aug_params: dict={'beta': 1.0}, soft_label_mode: str='S',
                  optimizer: str='sgd', lr_scheduler: str='step', temperature: float=1.0, weight_decay: float=0.0005, 
                  momentum: float=0.9, num_eval: int=5, im_size: tuple=(32, 32), num_epochs: int=300, use_zca: bool=False,
-                 real_batch_size: int=256, syn_batch_size: int=256, default_lr: float=0.01, save_path: str=None, 
+                 real_batch_size: int=256, syn_batch_size: int=256, default_lr: float=0.01, save_path: str=None, use_aug_for_hard: bool=False,
                  stu_use_torchvision: bool=False, tea_use_torchvision: bool=False, num_workers: int=4, teacher_dir: str='./teacher_models', 
                  custom_train_trans: transforms.Compose=None, custom_val_trans: transforms.Compose=None, device: str="cuda"):
 
@@ -60,6 +60,7 @@ def __init__(self, config: Config=None, dataset: str='CIFAR10', real_data_path:
                                                                                                    real_data_path, 
                                                                                                    im_size, 
                                                                                                    use_zca,
+                                                                                                   custom_train_trans,
                                                                                                    custom_val_trans,
                                                                                                    device)
         self.images_train, self.labels_train, self.class_indices_train = self.load_real_data(dst_train, class_map, num_classes)
@@ -99,6 +100,7 @@ def __init__(self, config: Config=None, dataset: str='CIFAR10', real_data_path:
             self.aug_func = Cutmix(aug_params)
         else:
             self.aug_func = None
+        self.use_aug_for_hard = use_aug_for_hard
 
         if not save_path:
             save_path = f"./results/{dataset}/{model_name}/ipc{ipc}/obj_scores.csv"
@@ -213,7 +215,7 @@ def compute_hard_label_metrics(self, model, image_tensor, image_path, lr, hard_l
                 loader=train_loader, 
                 loss_fn=loss_fn, 
                 optimizer=optimizer,
-                aug_func=self.aug_func,
+                aug_func=self.aug_func if self.use_aug_for_hard else None,
                 lr_scheduler=lr_scheduler, 
                 tea_model=self.teacher_model, 
                 device=self.device

dd_ranking/utils/data.py

@@ -48,7 +48,7 @@ def __len__(self):
         return len(self.images)
 
 
-def get_dataset(dataset, data_path, im_size, use_zca, custom_val_trans, device):
+def get_dataset(dataset, data_path, im_size, use_zca, custom_train_trans, custom_val_trans, device):
     class_map_inv = None
 
     if dataset == 'CIFAR10':
@@ -68,7 +68,7 @@ def get_dataset(dataset, data_path, im_size, use_zca, custom_val_trans, device):
                 transforms.ToTensor()
             ])
 
-        dst_train = datasets.CIFAR10(data_path, train=True, download=True, transform=transform)
+        dst_train = datasets.CIFAR10(data_path, train=True, download=True, transform=transform if custom_train_trans is None else custom_train_trans)
         dst_test = datasets.CIFAR10(data_path, train=False, download=True, transform=transform if custom_val_trans is None else custom_val_trans)
         class_map = {x: x for x in range(num_classes)}
 
@@ -89,7 +89,7 @@ def get_dataset(dataset, data_path, im_size, use_zca, custom_val_trans, device):
                 transforms.ToTensor()
             ])
 
-        dst_train = datasets.CIFAR100(data_path, train=True, download=True, transform=transform)
+        dst_train = datasets.CIFAR100(data_path, train=True, download=True, transform=transform if custom_train_trans is None else custom_train_trans)
         dst_test = datasets.CIFAR100(data_path, train=False, download=True, transform=transform if custom_val_trans is None else custom_val_trans)
         class_map = {x: x for x in range(num_classes)}
 
@@ -108,7 +108,7 @@ def get_dataset(dataset, data_path, im_size, use_zca, custom_val_trans, device):
             transform = transforms.Compose([
                 transforms.ToTensor()
             ])
-        dst_train = datasets.ImageFolder(os.path.join(data_path, "train"), transform=transform)
+        dst_train = datasets.ImageFolder(os.path.join(data_path, "train"), transform=transform if custom_train_trans is None else custom_train_trans)
         dst_test = datasets.ImageFolder(os.path.join(data_path, "val"), transform=transform if custom_val_trans is None else custom_val_trans)
         class_map = {x: x for x in range(num_classes)}
 
@@ -129,7 +129,7 @@ def get_dataset(dataset, data_path, im_size, use_zca, custom_val_trans, device):
             transforms.CenterCrop(im_size)
         ])
 
-        dst_train = datasets.ImageFolder(os.path.join(data_path, "train"), transform=transform)
+        dst_train = datasets.ImageFolder(os.path.join(data_path, "train"), transform=transform if custom_train_trans is None else custom_train_trans)
         dst_train = torch.utils.data.Subset(dst_train, np.squeeze(np.argwhere(np.isin(dst_train.targets, config.img_net_classes))))
         dst_test = datasets.ImageFolder(os.path.join(data_path, "val"), transform=transform if custom_val_trans is None else custom_val_trans)
         dst_test = torch.utils.data.Subset(dst_test, np.squeeze(np.argwhere(np.isin(dst_test.targets, config.img_net_classes))))
@@ -153,7 +153,7 @@ def get_dataset(dataset, data_path, im_size, use_zca, custom_val_trans, device):
             transforms.CenterCrop(im_size)
         ])
 
-        dst_train = datasets.ImageFolder(os.path.join(data_path, "train"), transform=transform)
+        dst_train = datasets.ImageFolder(os.path.join(data_path, "train"), transform=transform if custom_train_trans is None else custom_train_trans)
         dst_test = datasets.ImageFolder(os.path.join(data_path, "val"), transform=transform if custom_val_trans is None else custom_val_trans)
 
         class_map = {x: i for i, x in enumerate(range(num_classes))}