From 96ed9fe2aed6e670367d2e07d36a7f2d03c1eb9c Mon Sep 17 00:00:00 2001 From: "Harle, Antoine (Contracteur)" Date: Tue, 28 Jan 2020 19:42:00 -0500 Subject: [PATCH] Ajout simplification pour entrainment --- higher/smart_aug/dataug.py | 66 +++++++++++++++++++++++ higher/smart_aug/test_dataug.py | 1 + higher/smart_aug/train_utils.py | 96 +++++++++++++++++++++++++++++++++ 3 files changed, 163 insertions(+) diff --git a/higher/smart_aug/dataug.py b/higher/smart_aug/dataug.py index 1e7a100..688d8c8 100755 --- a/higher/smart_aug/dataug.py +++ b/higher/smart_aug/dataug.py @@ -870,6 +870,8 @@ class Higher_model(nn.Module): """ return self._name +from utils import clip_norm +from train_utils import compute_vaLoss class Augmented_model(nn.Module): """Wrapper for a Data Augmentation module and a model. @@ -917,6 +919,70 @@ class Augmented_model(nn.Module): self._data_augmentation=mode self._mods['data_aug'].augment(mode) + def start_bilevel_opt(self, inner_it, hp_list, opt_param, dl_val): + + if inner_it==0 or len(hp_list)==0: #No meta-opt + print("No meta optimization") + + self._diffopt = model['model'].get_diffopt( + inner_opt, + grad_callback=(lambda grads: clip_norm(grads, max_norm=10)), + track_higher_grads=False) + + else: #Bi-level opt + print("Bi-Level optimization") + self._it_count=0 + self._in_it=inner_it + + self._opt_param=opt_param + #Inner Opt + inner_opt = torch.optim.SGD(self._mods['model']['original'].parameters(), lr=opt_param['Inner']['lr'], momentum=opt_param['Inner']['momentum']) #lr=1e-2 / momentum=0.9 + + self._diffopt = self._mods['model'].get_diffopt( + inner_opt, + grad_callback=(lambda grads: clip_norm(grads, max_norm=10)), + track_higher_grads=True) + + #Meta Opt + self._meta_opt = torch.optim.Adam(hp_list, lr=opt_param['Meta']['lr']) + + self._dl_val=dl_val + self._dl_val_it=iter(dl_val) + self._val_loss=0. + + self._meta_opt.zero_grad() + + def step(self, loss): + + self._it_count+=1 + self._diffopt.step(loss) #(opt.zero_grad, loss.backward, opt.step) + + if(self._meta_opt and self._it_count>0 and self._it_count%self._in_it==0): #Perform Meta step + #print("meta") + self._val_loss = compute_vaLoss(model=self._mods['model'], dl_it=self._dl_val_it, dl=self._dl_val) + self._mods['data_aug'].reg_loss() + #print_graph(val_loss) #to visualize computational graph + self._val_loss.backward() + + torch.nn.utils.clip_grad_norm_(self._mods['data_aug'].parameters(), max_norm=10, norm_type=2) #Prevent exploding grad with RNN + + self._meta_opt.step() + + #Adjust Hyper-parameters + self._mods['data_aug'].adjust_param(soft=False) #Contrainte sum(proba)=1 + + #For optimizer parameters, if needed + #for param_group in self._diffopt.param_groups: + # for param in list(self._opt_param['Inner'].keys())[1:]: + # param_group[param].data = param_group[param].data.clamp(min=1e-4) + + #Reset gradients + self._diffopt.detach_() + self._mods['model'].detach_() + self._meta_opt.zero_grad() + + self._it_count=0 + + def train(self, mode=True): """ Set the module training mode. diff --git a/higher/smart_aug/test_dataug.py b/higher/smart_aug/test_dataug.py index 19da042..ce4c864 100755 --- a/higher/smart_aug/test_dataug.py +++ b/higher/smart_aug/test_dataug.py @@ -187,6 +187,7 @@ if __name__ == "__main__": #aug_model = Augmented_model(RandAug(TF_dict=tf_dict, N_TF=2), model).to(device) print("{} on {} for {} epochs - {} inner_it".format(str(aug_model), device_name, epochs, n_inner_iter)) + log= run_simple_smartaug(model=aug_model, opt_param=optim_param) log= run_dist_dataugV3(model=aug_model, epochs=epochs, inner_it=n_inner_iter, diff --git a/higher/smart_aug/train_utils.py b/higher/smart_aug/train_utils.py index 20e4ec8..c578034 100755 --- a/higher/smart_aug/train_utils.py +++ b/higher/smart_aug/train_utils.py @@ -363,3 +363,99 @@ def run_dist_dataugV3(model, opt_param, epochs=1, inner_it=1, dataug_epoch_start pass return log + +def run_simple_smartaug(model, opt_param, epochs=1, inner_it=1, print_freq=1, unsup_loss=1, save_sample_freq=None): + """Simple training of an augmented model with higher. + + This function is intended to be used with Augmented_model containing an Higher_model (see dataug.py). + Ex : Augmented_model(Data_augV5(...), Higher_model(model)) + + Training loss can either be computed directly from augmented inputs (unsup_loss=0). + However, it is recommended to use the mixed loss computation, which combine original and augmented inputs to compute the loss (unsup_loss>0). + + Args: + model (nn.Module): Augmented model to train. + opt_param (dict): Dictionnary containing optimizers parameters. + epochs (int): Number of epochs to perform. (default: 1) + inner_it (int): Number of inner iteration before a meta-step. 0 inner iteration means there's no meta-step. (default: 1) + print_freq (int): Number of epoch between display of the state of training. If set to None, no display will be done. (default:1) + unsup_loss (float): Proportion of the unsup_loss loss added to the supervised loss. If set to 0, the loss is only computed on augmented inputs. (default: 1) + save_sample_freq (int): Number of epochs between saves of samples of data. If set to None, only one save would be done at the end of the training. (default: None) + + Returns: + (list) Logs of training. Each items is a dict containing results of an epoch. + """ + device = next(model.parameters()).device + log = [] + + ## Optimizers ## + hyper_param = list(model['data_aug'].parameters()) + model.start_bilevel_opt(inner_it=inner_it, hp_list=hyper_param, opt_param=opt_param, dl_val=dl_val) + + model.train() + + for epoch in range(1, epochs+1): + t0 = time.process_time() + + for i, (xs, ys) in enumerate(dl_train): + xs, ys = xs.to(device), ys.to(device) + + #Methode mixed + loss = mixed_loss(xs, ys, model, unsup_factor=unsup_loss) + + model.step(loss) #(opt.zero_grad, loss.backward, opt.step) + automatic meta-optimisation + + tf = time.process_time() + + if (save_sample_freq and epoch%save_sample_freq==0): #Data sample saving + 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 samples epoch"+epoch) + pass + + val_loss = model._val_loss + # Test model + accuracy, test_loss =test(model) + model.train() + + #### Log #### + param = [{'p': p.item(), 'm':model['data_aug']['mag'].item()} for p in model['data_aug']['prob']] if model['data_aug']._shared_mag else [{'p': p.item(), 'm': m.item()} for p, m in zip(model['data_aug']['prob'], model['data_aug']['mag'])] + data={ + "epoch": epoch, + "train_loss": loss.item(), + "val_loss": val_loss.item(), + "acc": accuracy, + "time": tf - t0, + + "mix_dist": model['data_aug']['mix_dist'].item(), + "param": param, + } + log.append(data) + ############# + #### Print #### + if(print_freq and epoch%print_freq==0): + print('-'*9) + print('Epoch : %d/%d'%(epoch,epochs)) + print('Time : %.00f'%(tf - t0)) + print('Train loss :',loss.item(), '/ val loss', val_loss.item()) + print('Accuracy :', max([x["acc"] for x in log])) + print('Data Augmention : {} (Epoch {})'.format(model._data_augmentation, 0)) + if not model['data_aug']._fixed_prob: print('TF Proba :', model['data_aug']['prob'].data) + #print('proba grad',model['data_aug']['prob'].grad) + if not model['data_aug']._fixed_mag: print('TF Mag :', model['data_aug']['mag'].data) + #print('Mag grad',model['data_aug']['mag'].grad) + if not model['data_aug']._fixed_mix: print('Mix:', model['data_aug']['mix_dist'].item()) + #print('Reg loss:', model['data_aug'].reg_loss().item()) + ############# + + #Data sample saving + 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 + + return log