diff --git a/higher/compare_res.py b/higher/compare_res.py index fbe4949..c7ccd6e 100755 --- a/higher/compare_res.py +++ b/higher/compare_res.py @@ -82,7 +82,7 @@ if __name__ == "__main__": nb_run=3 accs = [] times = [] - files = ["res/brutus-tests/log/Aug_mod(Data_augV5(Uniform-14TFx2-MagFxSh)-LeNet)-150epochs(dataug:0)-0in_it-%s.json"%str(run) for run in range(nb_run)] + files = ["res/brutus-tests/log/Aug_mod(Data_augV5(Uniform-14TFx3-Mag)-LeNet)-150epochs(dataug:0)-1in_it-%s.json"%str(run) for run in range(nb_run)] for idx, file in enumerate(files): #legend+=str(idx)+'-'+file+'\n' diff --git a/higher/datasets.py b/higher/datasets.py index 88ec045..3688eda 100755 --- a/higher/datasets.py +++ b/higher/datasets.py @@ -50,7 +50,6 @@ class AugmentedDataset(VisionDataset): for idx, img in enumerate(self.sup_data): self.sup_data[idx]= Image.fromarray(img) #to PIL Image - self.unsup_ratio=5 #Batch size unsup = train batch size * unsup_ratio self.unsup_data=[] self.unsup_targets=[] @@ -157,6 +156,120 @@ class AugmentedDataset(VisionDataset): def __str__(self): return "CIFAR10(Sup:{}-Unsup:{}-{}TF)".format(self.dataset_info['sup'], self.dataset_info['unsup'], len(self._TF)) +class AugmentedDatasetV2(VisionDataset): + def __init__(self, root, train=True, transform=None, target_transform=None, download=False, subset=None): + + super(AugmentedDatasetV2, self).__init__(root, transform=transform, target_transform=target_transform) + + supervised_dataset = torchvision.datasets.CIFAR10(root, train=train, download=download, transform=transform) + + self.sup_data = supervised_dataset.data if not subset else supervised_dataset.data[subset[0]:subset[1]] + self.sup_targets = supervised_dataset.targets if not subset else supervised_dataset.targets[subset[0]:subset[1]] + assert len(self.sup_data)==len(self.sup_targets) + + for idx, img in enumerate(self.sup_data): + self.sup_data[idx]= Image.fromarray(img) #to PIL Image + + self.unsup_data=[] + self.unsup_targets=[] + self.origin_idx=[] + + self.dataset_info= { + 'name': 'CIFAR10', + 'sup': len(self.sup_data), + 'unsup': len(self.unsup_data), + 'length': len(self.sup_data)+len(self.unsup_data), + } + + + self._TF = [ + ## Geometric TF ## + 'Rotate', + 'TranslateX', + 'TranslateY', + 'ShearX', + 'ShearY', + + 'Cutout', + + ## Color TF ## + 'Contrast', + 'Color', + 'Brightness', + 'Sharpness', + #'Posterize', + #'Solarize', + + 'Invert', + 'AutoContrast', + 'Equalize', + ] + self._op_list =[] + self.prob=0.5 + for tf in self._TF: + for mag in range(1, 10): + self._op_list+=[(tf, self.prob, mag)] + self._nb_op = len(self._op_list) + + def __getitem__(self, index): + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + aug_img, origin_img, target = self.unsup_data[index], self.sup_data[self.origin_idx[index]], self.unsup_targets[index] + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + #img = Image.fromarray(img) + + if self.transform is not None: + aug_img = self.transform(aug_img) + origin_img = self.transform(origin_img) + + if self.target_transform is not None: + target = self.target_transform(target) + + return aug_img, origin_img, target + + def augement_data(self, aug_copy=1): + + policies = [] + for op_1 in self._op_list: + for op_2 in self._op_list: + policies += [[op_1, op_2]] + + for idx, image in enumerate(self.sup_data): + if idx%(self.dataset_info['sup']/5)==0: print("Augmenting data... ", idx,"/", self.dataset_info['sup']) + #if idx==10000:break + + for _ in range(aug_copy): + chosen_policy = policies[np.random.choice(len(policies))] + aug_image = augmentation_transforms.apply_policy(chosen_policy, image, use_mean_std=False) #Cast en float image + #aug_image = augmentation_transforms.cutout_numpy(aug_image) + + self.unsup_data+=[(aug_image*255.).astype(self.sup_data.dtype)]#Cast float image to uint8 + self.unsup_targets+=[self.sup_targets[idx]] + self.origin_idx+=[idx] + + #self.unsup_data=(np.array(self.unsup_data)*255.).astype(self.sup_data.dtype) #Cast float image to uint8 + self.unsup_data=np.array(self.unsup_data) + + assert len(self.unsup_data)==len(self.unsup_targets) + + self.dataset_info['unsup']=len(self.unsup_data) + self.dataset_info['length']=self.dataset_info['sup']+self.dataset_info['unsup'] + + + def __len__(self): + return self.dataset_info['unsup']#self.dataset_info['length'] + + def __str__(self): + return "CIFAR10(Sup:{}-Unsup:{}-{}TF)".format(self.dataset_info['sup'], self.dataset_info['unsup'], len(self._TF)) + + ### Classic Dataset ### data_train = torchvision.datasets.CIFAR10("./data", train=True, download=download_data, transform=transform) #data_val = torchvision.datasets.CIFAR10("./data", train=True, download=download_data, transform=transform) diff --git a/higher/test_brutus.py b/higher/test_brutus.py index 8a5dd85..88de3ca 100755 --- a/higher/test_brutus.py +++ b/higher/test_brutus.py @@ -35,7 +35,7 @@ if __name__ == "__main__": n_inner_iter = 1 - epochs = 150 + epochs = 200 dataug_epoch_start=0 #model = LeNet(3,10) @@ -44,22 +44,6 @@ if __name__ == "__main__": tf_dict = {k: TF.TF_dict[k] for k in tf_names} - t0 = time.process_time() - - aug_model = Augmented_model(Data_augV5(TF_dict=tf_dict, N_TF=3, mix_dist=0.0, fixed_prob=True, fixed_mag=True, shared_mag=True), model).to(device) - - print("{} on {} for {} epochs - {} inner_it".format(str(aug_model), device_name, epochs, n_inner_iter)) - log= run_dist_dataugV2(model=aug_model, epochs=epochs, inner_it=n_inner_iter, dataug_epoch_start=dataug_epoch_start, print_freq=10, KLdiv=True, loss_patience=None) - - exec_time=time.process_time() - t0 - #### - times = [x["time"] for x in log] - out = {"Accuracy": max([x["acc"] for x in log]), "Time": (np.mean(times),np.std(times), exec_time), "Device": device_name, "Param_names": aug_model.TF_names(), "Log": log} - filename = "{}-{} epochs (dataug:{})- {} in_it".format(str(aug_model),epochs,dataug_epoch_start,n_inner_iter) - with open("res/log/%s.json" % filename, "w+") as f: - json.dump(out, f, indent=True) - print('Log :\"',f.name, '\" saved !') - #### t0 = time.process_time() @@ -77,6 +61,23 @@ if __name__ == "__main__": json.dump(out, f, indent=True) print('Log :\"',f.name, '\" saved !') + + #### + t0 = time.process_time() + + aug_model = Augmented_model(Data_augV5(TF_dict=tf_dict, N_TF=3, mix_dist=0.0, fixed_prob=False, fixed_mag=False, shared_mag=False), model).to(device) + + print("{} on {} for {} epochs - {} inner_it".format(str(aug_model), device_name, epochs, n_inner_iter)) + log= run_dist_dataugV2(model=aug_model, epochs=epochs, inner_it=n_inner_iter, dataug_epoch_start=dataug_epoch_start, print_freq=10, KLdiv=True, loss_patience=None) + + exec_time=time.process_time() - t0 + #### + times = [x["time"] for x in log] + out = {"Accuracy": max([x["acc"] for x in log]), "Time": (np.mean(times),np.std(times), exec_time), "Device": device_name, "Param_names": aug_model.TF_names(), "Log": log} + filename = "{}-{} epochs (dataug:{})- {} in_it".format(str(aug_model),epochs,dataug_epoch_start,n_inner_iter) + with open("res/log/%s.json" % filename, "w+") as f: + json.dump(out, f, indent=True) + print('Log :\"',f.name, '\" saved !') ''' res_folder="res/brutus-tests/" epochs= 150 diff --git a/higher/train_utils.py b/higher/train_utils.py index 856d884..c1fc880 100755 --- a/higher/train_utils.py +++ b/higher/train_utils.py @@ -629,7 +629,7 @@ def run_dist_dataug(model, epochs=1, inner_it=1, dataug_epoch_start=0): print("Copy ", countcopy) return log -def run_dist_dataugV2(model, epochs=1, inner_it=0, dataug_epoch_start=0, print_freq=1, KLdiv=False, loss_patience=None): +def run_dist_dataugV2(model, epochs=1, inner_it=0, dataug_epoch_start=0, print_freq=1, KLdiv=False, loss_patience=None, save_sample=False): device = next(model.parameters()).device log = [] countcopy=0 @@ -796,8 +796,12 @@ def run_dist_dataugV2(model, epochs=1, inner_it=0, dataug_epoch_start=0, print_f model.augment(mode=True) if inner_it != 0: high_grad_track = True - viz_sample_data(imgs=xs, labels=ys, fig_name='samples/data_sample_epoch{}_noTF'.format(epoch)) - viz_sample_data(imgs=model['data_aug'](xs), labels=ys, fig_name='samples/data_sample_epoch{}'.format(epoch)) + try: + viz_sample_data(imgs=xs, labels=ys, fig_name='samples/data_sample_epoch{}_noTF'.format(epoch)) + viz_sample_data(imgs=model['data_aug'](xs), labels=ys, fig_name='samples/data_sample_epoch{}'.format(epoch)) + except: + print("Couldn't save finals samples") + pass #print("Copy ", countcopy) return log