class MAMLFewShotClassifier(nn.Module): def __init__(self, im_shape, device, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLFewShotClassifier, self).__init__() self.args = args self.device = device self.batch_size = args.batch_size self.use_cuda = args.use_cuda self.im_shape = im_shape self.current_epoch = 0 self.rng = set_torch_seed(seed=args.seed) if not self.args.retrain: print('search for optimal arch for meta learning') self.classifier = MetaPruningNetwork( args=args, in_channels=im_shape[1], init_channels=args.init_channels, num_classes=self.args.num_classes_per_set, layers=self.args.layers, device=device).to(device=self.device) else: # for retrain print('retraining the arch', args.arch, eval('meta_genotype.{}'.format(args.arch))) self.classifier = MetaPrunedPretrainedNetwork( args=args, genotype=eval('meta_genotype.{}'.format(args.arch)), in_channels=im_shape[1], init_channels=args.init_channels, num_classes=self.args.num_classes_per_set, device=device).to(device=self.device) # self.task_learning_rate = args.task_learning_rate #0.1 # self.task_learning_rate = args.init_inner_loop_learning_rate #0.1 self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=device, init_learning_rate=args.init_inner_loop_learning_rate, init_learning_rate_arch=args.init_inner_arch_loop_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) if self.args.retrain: # only weights self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters(), )) else: # arch + weights self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters(), arch_params=self.classifier.named_arch_parameters())) self.use_cuda = args.use_cuda self.device = device self.args = args self.to(device) # print("Inner Loop parameters") # for key, value in self.inner_loop_optimizer.named_parameters(): # print(key, value.shape) # print("Outer Loop parameters") # for name, param in self.named_parameters(): # if param.requires_grad: # print(name, param.shape, param.device, param.requires_grad) self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.meta_learning_rate, amsgrad=False) if not self.args.retrain: # train the arch self.optimizer_arch = optim.Adam( self.classifier.arch_parameters(), lr=args.init_meta_arch_loop_learning_rate, betas=(0.5, 0.999), weight_decay=args.meta_arch_weights_decay) self.scheduler_arch = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer_arch, T_max=self.args.total_epochs, eta_min=self.args.min_arch_learning_rate) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate) self.device = torch.device('cpu') if torch.cuda.is_available(): # print('torch.cuda.device_count() = ', torch.cuda.device_count()) if torch.cuda.device_count() > 1: self.to(torch.cuda.current_device()) self.classifier = nn.DataParallel(module=self.classifier) else: self.to(torch.cuda.current_device()) self.device = torch.cuda.current_device() def get_per_step_loss_importance_vector(self): """ Generates a tensor of dimensionality (num_inner_loop_steps) indicating the importance of each step's target loss towards the optimization loss. :return: A tensor to be used to compute the weighted average of the loss, useful for the MSL (Multi Step Loss) mechanism. """ loss_weights = np.ones( shape=(self.args.number_of_training_steps_per_iter)) * ( 1.0 / self.args.number_of_training_steps_per_iter) decay_rate = 1.0 / self.args.number_of_training_steps_per_iter / self.args.multi_step_loss_num_epochs min_value_for_non_final_losses = 0.03 / self.args.number_of_training_steps_per_iter for i in range(len(loss_weights) - 1): curr_value = np.maximum( loss_weights[i] - (self.current_epoch * decay_rate), min_value_for_non_final_losses) loss_weights[i] = curr_value curr_value = np.minimum( loss_weights[-1] + (self.current_epoch * (self.args.number_of_training_steps_per_iter - 1) * decay_rate), 1.0 - ((self.args.number_of_training_steps_per_iter - 1) * min_value_for_non_final_losses)) loss_weights[-1] = curr_value loss_weights = torch.Tensor(loss_weights).to(device=self.device) return loss_weights def get_inner_loop_parameter_dict(self, params=None, arch_params=None): """ Returns a dictionary with the parameters to use for inner loop updates. :param params: A dictionary of the network's parameters. :return: A dictionary of the parameters to use for the inner loop optimization process. """ param_dict = dict() if params is not None: for name, param in params: if param.requires_grad: if self.args.enable_inner_loop_optimizable_bn_params: # enable bn params to be optimized param_dict[name] = param.to(device=self.device) # param_dict[name] = param else: if "norm_layer" not in name: param_dict[name] = param.to(device=self.device) # param_dict[name] = param if arch_params is not None: for name, param in arch_params: if param.requires_grad: param_dict[name] = param.to(device=self.device) return param_dict def apply_inner_loop_update(self, loss, names_weights_copy, use_second_order, current_step_idx): """ Applies an inner loop update given current step's loss, the weights to update, a flag indicating whether to use second order derivatives and the current step's index. :param loss: Current step's loss with respect to the support set. :param names_weights_copy: A dictionary with names to parameters to update. :param use_second_order: A boolean flag of whether to use second order derivatives. :param current_step_idx: Current step's index. :return: A dictionary with the updated weights (name, param) """ num_gpus = torch.cuda.device_count() if num_gpus > 1: self.classifier.module.zero_grad(params=names_weights_copy) else: self.classifier.zero_grad(params=names_weights_copy) grads = torch.autograd.grad(loss, names_weights_copy.values(), create_graph=use_second_order, allow_unused=True) names_grads_copy = dict(zip(names_weights_copy.keys(), grads)) names_weights_copy = { key: value[0] for key, value in names_weights_copy.items() } for key, grad in names_grads_copy.items(): if grad is None: # some grads are not needed. raise Exception('Grads not found for inner loop parameter', key) else: names_grads_copy[key] = names_grads_copy[key].sum(dim=0) names_weights_copy = self.inner_loop_optimizer.update_params( names_weights_dict=names_weights_copy, names_grads_wrt_params_dict=names_grads_copy, num_step=current_step_idx) num_devices = torch.cuda.device_count() if torch.cuda.is_available( ) else 1 names_weights_copy = { name.replace('module.', ''): value.unsqueeze(0).repeat([num_devices] + [1 for i in range(len(value.shape))]) for name, value in names_weights_copy.items() } return names_weights_copy def get_across_task_loss_metrics(self, total_losses, total_accuracies): losses = dict() losses['loss'] = torch.mean( torch.stack(total_losses)) #use loss average losses['accuracy'] = np.mean(total_accuracies) return losses def forward(self, data_batch, epoch, use_second_order, use_multi_step_loss_optimization, num_steps, training_phase, param='weight'): """ Runs a forward outer loop pass on the batch of tasks using the MAML/++ framework. :param data_batch: A data batch containing the support and target sets. :param epoch: Current epoch's index :param use_second_order: A boolean saying whether to use second order derivatives. :param use_multi_step_loss_optimization: Whether to optimize on the outer loop using just the last step's target loss (True) or whether to use multi step loss which improves the stability of the system (False) :param num_steps: Number of inner loop steps. :param training_phase: Whether this is a training phase (True) or an evaluation phase (False) :param param: 0 for weights. 1 for alpha :return: A dictionary with the collected losses of the current outer forward propagation. """ x_support_set, x_target_set, y_support_set, y_target_set = data_batch total_losses = [] total_accuracies = [] per_task_target_preds = [[] for i in range(len(x_target_set))] self.classifier.zero_grad() for task_id, (x_support_set_task, y_support_set_task, x_target_set_task, y_target_set_task) in \ enumerate(zip(x_support_set, y_support_set, x_target_set, y_target_set)): task_losses = [] task_accuracies = [] per_step_loss_importance_vectors = self.get_per_step_loss_importance_vector( ) names_weights_copy = None if param == 'weight': names_weights_copy = self.get_inner_loop_parameter_dict( self.classifier.named_parameters(), ) else: # for alpha names_weights_copy = self.get_inner_loop_parameter_dict( arch_params=self.classifier.named_arch_parameters(), ) num_devices = torch.cuda.device_count() if torch.cuda.is_available( ) else 1 names_weights_copy = { name.replace('module.', ''): value.unsqueeze(0).repeat([num_devices] + [1 for i in range(len(value.shape))]) for name, value in names_weights_copy.items() } n, s, c, h, w = x_target_set_task.shape x_support_set_task = x_support_set_task.view(-1, c, h, w) y_support_set_task = y_support_set_task.view(-1) x_target_set_task = x_target_set_task.view(-1, c, h, w) y_target_set_task = y_target_set_task.view(-1) task_max_acc = 0.0 for num_step in range(num_steps): support_loss, support_preds = self.net_forward( x=x_support_set_task, y=y_support_set_task, weights=names_weights_copy, backup_running_statistics=True if (num_step == 0) else False, training=True, num_step=num_step) names_weights_copy = self.apply_inner_loop_update( loss=support_loss, names_weights_copy=names_weights_copy, use_second_order=use_second_order, current_step_idx=num_step) if use_multi_step_loss_optimization and training_phase and epoch < self.args.multi_step_loss_num_epochs: target_loss, target_preds = self.net_forward( x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) task_losses.append( per_step_loss_importance_vectors[num_step] * target_loss) else: if num_step == ( self.args.number_of_training_steps_per_iter - 1): target_loss, target_preds = self.net_forward( x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) task_losses.append(target_loss) per_task_target_preds[task_id] = target_preds.detach().cpu().numpy( ) _, predicted = torch.max(target_preds.data, 1) accuracy = predicted.float().eq( y_target_set_task.data.float()).cpu().float() task_losses = torch.sum(torch.stack(task_losses)) total_losses.append(task_losses) total_accuracies.extend(accuracy) if not training_phase: self.classifier.restore_backup_stats() losses = self.get_across_task_loss_metrics( total_losses=total_losses, total_accuracies=total_accuracies) # for idx, item in enumerate(per_step_loss_importance_vectors): # losses['loss_importance_vector_{}'.format(idx)] = item.detach().cpu().numpy() return losses, per_task_target_preds def net_forward(self, x, y, weights, backup_running_statistics, training, num_step): """ A base model forward pass on some data points x. Using the parameters in the weights dictionary. Also requires boolean flags indicating whether to reset the running statistics at the end of the run (if at evaluation phase). A flag indicating whether this is the training session and an int indicating the current step's number in the inner loop. :param x: A data batch of shape b, c, h, w :param y: A data targets batch of shape b, n_classes :param weights: A dictionary containing the weights to pass to the network. :param backup_running_statistics: A flag indicating whether to reset the batch norm running statistics to their previous values after the run (only for evaluation) :param training: A flag indicating whether the current process phase is a training or evaluation. :param num_step: An integer indicating the number of the step in the inner loop. :return: the crossentropy losses with respect to the given y, the predictions of the base model. """ preds = self.classifier.forward( x=x, params=weights, training=training, backup_running_statistics=backup_running_statistics, num_step=num_step, ) loss = F.cross_entropy(input=preds, target=y) return loss, preds def trainable_parameters(self): """ Returns an iterator over the trainable parameters of the model. """ for param in self.parameters(): if param.requires_grad: yield param def train_forward_prop(self, data_batch, epoch, param='weight'): """ Runs an outer loop forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, per_task_target_preds = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=self.args.second_order and epoch > self.args.first_order_to_second_order_epoch, use_multi_step_loss_optimization=self.args. use_multi_step_loss_optimization, num_steps=self.args.number_of_training_steps_per_iter, training_phase=True, param=param) return losses, per_task_target_preds def train_arch_forward_prop(self, data_batch, epoch, param='arch'): """ Runs an outer loop forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, per_task_target_preds = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=False, use_multi_step_loss_optimization=self.args. use_multi_step_loss_optimization, num_steps=self.args.number_of_training_steps_per_iter, training_phase=True, param=param) return losses, per_task_target_preds def evaluation_forward_prop(self, data_batch, epoch): """ Runs an outer loop evaluation forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, per_task_target_preds = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=False, use_multi_step_loss_optimization=True, num_steps=self.args.number_of_evaluation_steps_per_iter, training_phase=False) return losses, per_task_target_preds def meta_update_weights(self, loss): """ Applies an outer loop update on the meta-parameters of the model. :param loss: The current crossentropy loss. """ self.optimizer.zero_grad() # self.optimizer_arch.zero_grad() loss.backward() if 'imagenet' in self.args.dataset_name: for name, param in self.classifier.named_parameters(): if param.requires_grad: param.grad.data.clamp_( -10, 10 ) # not sure if this is necessary, more experiments are needed self.optimizer.step() # self.optimizer_arch.step() def meta_update_alphas(self, loss): """ Applies an outer loop update on the meta-parameters of the model. :param loss: The current crossentropy loss. """ self.optimizer_arch.zero_grad() loss.backward() self.optimizer_arch.step() def batch_helper(self, data_batch, split=False): x_support_set, x_target_set, y_support_set, y_target_set = data_batch x_support_set = torch.Tensor(x_support_set).float().to( device=self.device) x_target_set = torch.Tensor(x_target_set).float().to( device=self.device) y_support_set = torch.Tensor(y_support_set).long().to( device=self.device) y_target_set = torch.Tensor(y_target_set).long().to(device=self.device) if not split: data_batch = (x_support_set, x_target_set, y_support_set, y_target_set) return data_batch else: #split the data task_num_split = len(x_target_set) // 2 data_batch_arch = (x_support_set[:task_num_split], x_target_set[:task_num_split], y_support_set[:task_num_split], y_target_set[:task_num_split]) data_batch_weight = (x_support_set[task_num_split:], x_target_set[task_num_split:], y_support_set[task_num_split:], y_target_set[task_num_split:]) return data_batch_arch, data_batch_weight def run_train_iter(self, data_batch, sample_idx, epoch): """ Runs an outer loop update step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ epoch = int(epoch) # self.step = sample_idx self.scheduler.step(epoch=epoch) if not self.args.retrain: self.scheduler_arch.step(epoch=epoch) if self.current_epoch != epoch: self.current_epoch = epoch # print('learning_rate = ', self.scheduler.get_last_lr()[0]) # self.scheduler.step() if not self.training: self.train() if not self.args.retrain: # search process data_batch_arch, data_batch_weight = self.batch_helper(data_batch, split=True) train_prop_weight = self.train_arch_forward_prop if epoch > self.args.min_search_epoch: if torch.sum(self.classifier.normal_candidate_flags.int()) > 0: losses, _ = self.train_arch_forward_prop( data_batch=data_batch_arch, epoch=epoch, param='arch') self.meta_update_alphas(loss=losses['loss']) self.optimizer.zero_grad() self.optimizer_arch.zero_grad() self.zero_grad() else: # no need for optimize arch train_prop_weight = self.train_forward_prop losses, per_task_target_preds = train_prop_weight( data_batch=data_batch_weight, epoch=epoch, param='weight') self.meta_update_weights(loss=losses['loss']) losses['learning_rate'] = self.scheduler.get_last_lr()[0] self.optimizer.zero_grad() self.optimizer_arch.zero_grad() self.zero_grad() else: # different batch data_batch = self.batch_helper(data_batch) losses, per_task_target_preds = self.train_forward_prop( data_batch=data_batch, epoch=epoch, param='weight') self.meta_update_weights(loss=losses['loss']) losses['learning_rate'] = self.scheduler.get_last_lr()[0] # print("current_learning_rate =", losses['learning_rate']) self.optimizer.zero_grad() self.zero_grad() return losses, per_task_target_preds def run_test_iter(self, data_batch): """ Runs an outer loop evaluation step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ if self.training: self.eval() x_support_set, x_target_set, y_support_set, y_target_set = data_batch x_support_set = torch.Tensor(x_support_set).float().to( device=self.device) x_target_set = torch.Tensor(x_target_set).float().to( device=self.device) y_support_set = torch.Tensor(y_support_set).long().to( device=self.device) y_target_set = torch.Tensor(y_target_set).long().to(device=self.device) data_batch = (x_support_set, x_target_set, y_support_set, y_target_set) losses, per_task_target_preds = self.evaluation_forward_prop( data_batch=data_batch, epoch=self.current_epoch) # losses['loss'].backward() # uncomment if you get the weird memory error # self.zero_grad() # self.optimizer.zero_grad() return losses, per_task_target_preds def run_check_iter(self, data_batch, logging): """ check weather all the weights still require grad. """ x_support_set, x_target_set, y_support_set, y_target_set = data_batch x_support_set = torch.Tensor(x_support_set).float().to( device=self.device) x_target_set = torch.Tensor(x_target_set).float().to( device=self.device) y_support_set = torch.Tensor(y_support_set).long().to( device=self.device) y_target_set = torch.Tensor(y_target_set).long().to(device=self.device) if not self.training: self.train() for task_id, (x_support_set_task, y_support_set_task, x_target_set_task, y_target_set_task) in \ enumerate(zip(x_support_set, y_support_set, x_target_set, y_target_set)): names_weights_copy = self.get_inner_loop_parameter_dict( self.classifier.named_parameters()) num_devices = torch.cuda.device_count() if torch.cuda.is_available( ) else 1 names_weights_copy = { name.replace('module.', ''): value.unsqueeze(0).repeat([num_devices] + [1 for i in range(len(value.shape))]) for name, value in names_weights_copy.items() } n, s, c, h, w = x_target_set_task.shape x_support_set_task = x_support_set_task.view(-1, c, h, w) y_support_set_task = y_support_set_task.view(-1) x_target_set_task = x_target_set_task.view(-1, c, h, w) y_target_set_task = y_target_set_task.view(-1) num_step = 0 target_loss, target_preds = self.net_forward( x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) self.optimizer.zero_grad() target_loss.backward() for name, param in self.classifier.named_parameters(): # for name, param in self.named_parameters(): if param.requires_grad: if param.grad is not None: # print(name, param.grad.shape) if torch.sum(param.grad) == 0: if 'op' in name: logging.info( '{} requires_grad = False'.format(name)) param.requires_grad = False else: if 'op' in name: logging.info( '{} requires_grad = False'.format(name)) param.requires_grad = False self.optimizer.zero_grad() self.zero_grad() break # only one iter def save_model(self, model_save_dir, state): """ Save the network parameter state and experiment state dictionary. :param model_save_dir: The directory to store the state at. :param state: The state containing the experiment state and the network. It's in the form of a dictionary object. """ state['network'] = self.state_dict() state['classifier'] = self.classifier.state_dict() torch.save(state, f=model_save_dir) def load_model(self, model_save_dir, model_name, model_idx): """ Load checkpoint and return the state dictionary containing the network state params and experiment state. :param model_save_dir: The directory from which to load the files. :param model_name: The model_name to be loaded from the direcotry. :param model_idx: The index of the model (i.e. epoch number or 'latest' for the latest saved model of the current experiment) :return: A dictionary containing the experiment state and the saved model parameters. """ filepath = os.path.join(model_save_dir, "{}_{}".format(model_name, model_idx)) state = torch.load(filepath) state_dict_loaded = state['network'] if not self.args.retrain: self.classifier.load_alphas(state) self.load_state_dict(state_dict=state_dict_loaded) return state def load_pretrained_model_init(self, pretrained_mode_path): state = torch.load(pretrained_mode_path) state_dict_loaded = state['network'] state_dict_loaded = {k[len('classifier.'):]: v for k, v in state_dict_loaded.items() \ if k.startswith('classifier.') and k[len('classifier.'):] in self.classifier.state_dict()} # pretrained_model = dict() for k, v in state_dict_loaded.items(): if self.args.per_step_bn_statistics: if 'norm' in k and len(v.shape) == 1: # print(k) state_dict_loaded[k] = v.unsqueeze(0).repeat( [self.args.number_of_training_steps_per_iter] + [1 for i in range(len(v.shape))]) else: if 'norm' in k and len(v.shape) == 2: # unuse the per_step_bn_statistics state_dict_loaded[k] = v[0] self.classifier.load_state_dict(state_dict_loaded)
def __init__(self, im_shape, device, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLFewShotClassifier, self).__init__() self.args = args self.device = device self.batch_size = args.batch_size self.use_cuda = args.use_cuda self.im_shape = im_shape self.current_epoch = 0 self.rng = set_torch_seed(seed=args.seed) if not self.args.retrain: print('search for optimal arch for meta learning') self.classifier = MetaPruningNetwork( args=args, in_channels=im_shape[1], init_channels=args.init_channels, num_classes=self.args.num_classes_per_set, layers=self.args.layers, device=device).to(device=self.device) else: # for retrain print('retraining the arch', args.arch, eval('meta_genotype.{}'.format(args.arch))) self.classifier = MetaPrunedPretrainedNetwork( args=args, genotype=eval('meta_genotype.{}'.format(args.arch)), in_channels=im_shape[1], init_channels=args.init_channels, num_classes=self.args.num_classes_per_set, device=device).to(device=self.device) # self.task_learning_rate = args.task_learning_rate #0.1 # self.task_learning_rate = args.init_inner_loop_learning_rate #0.1 self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=device, init_learning_rate=args.init_inner_loop_learning_rate, init_learning_rate_arch=args.init_inner_arch_loop_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) if self.args.retrain: # only weights self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters(), )) else: # arch + weights self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters(), arch_params=self.classifier.named_arch_parameters())) self.use_cuda = args.use_cuda self.device = device self.args = args self.to(device) # print("Inner Loop parameters") # for key, value in self.inner_loop_optimizer.named_parameters(): # print(key, value.shape) # print("Outer Loop parameters") # for name, param in self.named_parameters(): # if param.requires_grad: # print(name, param.shape, param.device, param.requires_grad) self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.meta_learning_rate, amsgrad=False) if not self.args.retrain: # train the arch self.optimizer_arch = optim.Adam( self.classifier.arch_parameters(), lr=args.init_meta_arch_loop_learning_rate, betas=(0.5, 0.999), weight_decay=args.meta_arch_weights_decay) self.scheduler_arch = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer_arch, T_max=self.args.total_epochs, eta_min=self.args.min_arch_learning_rate) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate) self.device = torch.device('cpu') if torch.cuda.is_available(): # print('torch.cuda.device_count() = ', torch.cuda.device_count()) if torch.cuda.device_count() > 1: self.to(torch.cuda.current_device()) self.classifier = nn.DataParallel(module=self.classifier) else: self.to(torch.cuda.current_device()) self.device = torch.cuda.current_device()
class MAMLFewShotClassifier(nn.Module): def __init__(self, im_shape, device, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLFewShotClassifier, self).__init__() self.args = args self.device = device self.batch_size = args.batch_size self.use_cuda = args.use_cuda self.im_shape = im_shape self.current_epoch = 0 self.rng = set_torch_seed(seed=args.seed) if args.high_end: self.embedding = HighEndEmbedding(device, args, 3).to(device) self.classifier = HighEndClassifier( device, args, self.embedding.n_out_channels).to(device) else: self.classifier = VGGReLUNormNetwork( im_shape=self.im_shape, num_output_classes=self.args.num_classes_per_set, args=args, device=device, meta_classifier=True).to(device=self.device) self.task_learning_rate = args.task_learning_rate self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=device, init_learning_rate=self.task_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter + self.args.num_critic_updates, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters())) print("Inner Loop parameters") for key, value in self.inner_loop_optimizer.named_parameters(): print(key, value.shape) if args.use_critic: print( sum([ reduce(mul, p.size(), 1) for p in list( self.get_inner_loop_parameter_dict( self.classifier.named_parameters()).values()) ])) self.critic = Critic(n_theta=sum([ reduce(mul, p.size(), 1) for p in list( self.get_inner_loop_parameter_dict( self.classifier.named_parameters()).values()) ])) self.use_cuda = args.use_cuda self.device = device self.args = args self.to(device) print("Outer Loop parameters") for name, param in self.named_parameters(): if param.requires_grad: print(name, param.shape, param.device, param.requires_grad) if args.high_end: self.optimizer = optim.SGD(self.trainable_parameters(), lr=1e-4) else: self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.meta_learning_rate, amsgrad=False) if args.use_critic: self.critic_optimizer = optim.SGD(self.critic.parameters(), lr=1e-6) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate) def get_per_step_loss_importance_vector(self): """ Generates a tensor of dimensionality (num_inner_loop_steps) indicating the importance of each step's target loss towards the optimization loss. :return: A tensor to be used to compute the weighted average of the loss, useful for the MSL (Multi Step Loss) mechanism. """ loss_weights = np.ones( shape=(self.args.number_of_training_steps_per_iter)) * ( 1.0 / self.args.number_of_training_steps_per_iter) decay_rate = 1.0 / self.args.number_of_training_steps_per_iter / self.args.multi_step_loss_num_epochs min_value_for_non_final_losses = 0.03 / self.args.number_of_training_steps_per_iter for i in range(len(loss_weights) - 1): curr_value = np.maximum( loss_weights[i] - (self.current_epoch * decay_rate), min_value_for_non_final_losses) loss_weights[i] = curr_value curr_value = np.minimum( loss_weights[-1] + (self.current_epoch * (self.args.number_of_training_steps_per_iter - 1) * decay_rate), 1.0 - ((self.args.number_of_training_steps_per_iter - 1) * min_value_for_non_final_losses)) loss_weights[-1] = curr_value loss_weights = torch.Tensor(loss_weights).to(device=self.device) return loss_weights def get_inner_loop_parameter_dict(self, params): """ Returns a dictionary with the parameters to use for inner loop updates. :param params: A dictionary of the network's parameters. :return: A dictionary of the parameters to use for the inner loop optimization process. """ param_dict = dict() for name, param in params: if param.requires_grad: if self.args.enable_inner_loop_optimizable_bn_params: param_dict[name] = param.to(device=self.device) else: if "norm_layer" not in name: param_dict[name] = param.to(device=self.device) return param_dict def apply_inner_loop_update(self, loss, names_weights_copy, use_second_order, current_step_idx): """ Applies an inner loop update given current step's loss, the weights to update, a flag indicating whether to use second order derivatives and the current step's index. :param loss: Current step's loss with respect to the support set. :param names_weights_copy: A dictionary with names to parameters to update. :param use_second_order: A boolean flag of whether to use second order derivatives. :param current_step_idx: Current step's index. :return: A dictionary with the updated weights (name, param) """ self.classifier.zero_grad(names_weights_copy) grads = torch.autograd.grad(loss, names_weights_copy.values(), create_graph=use_second_order) names_grads_wrt_params = dict(zip(names_weights_copy.keys(), grads)) names_weights_copy = self.inner_loop_optimizer.update_params( names_weights_dict=names_weights_copy, names_grads_wrt_params_dict=names_grads_wrt_params, num_step=current_step_idx) return names_weights_copy def get_across_task_loss_metrics(self, total_losses, total_accuracies): losses = dict() losses['loss'] = torch.mean(torch.stack(total_losses)) losses['accuracy'] = np.mean(total_accuracies) return losses def forward(self, data_batch, epoch, use_second_order, use_multi_step_loss_optimization, num_steps, training_phase): """ Runs a forward outer loop pass on the batch of tasks using the MAML/++ framework. :param data_batch: A data batch containing the support and target sets. :param epoch: Current epoch's index :param use_second_order: A boolean saying whether to use second order derivatives. :param use_multi_step_loss_optimization: Whether to optimize on the outer loop using just the last step's target loss (True) or whether to use multi step loss which improves the stability of the system (False) :param num_steps: Number of inner loop steps. :param training_phase: Whether this is a training phase (True) or an evaluation phase (False) :return: A dictionary with the collected losses of the current outer forward propagation. """ x_support_set, x_target_set, y_support_set, y_target_set = data_batch [b, ncs, spc] = y_support_set.shape self.num_classes_per_set = ncs total_losses = [] total_accuracies = [] per_task_target_preds = [[] for i in range(len(x_target_set))] if self.args.high_end: self.embedding.zero_grad() self.classifier.zero_grad() for task_id, (x_support_set_task, y_support_set_task, x_target_set_task, y_target_set_task) in \ enumerate(zip(x_support_set, y_support_set, x_target_set, y_target_set)): task_losses = [] task_accuracies = [] per_step_loss_importance_vectors = self.get_per_step_loss_importance_vector( ) # this is theta_0 names_weights_copy = self.get_inner_loop_parameter_dict( self.classifier.named_parameters()) n, s, c, h, w = x_target_set_task.shape x_support_set_task = x_support_set_task.view(-1, c, h, w) y_support_set_task = y_support_set_task.view(-1) x_target_set_task = x_target_set_task.view(-1, c, h, w) y_target_set_task = y_target_set_task.view(-1) # Inner loop starts if self.args.high_end: x_support_set_task = self.embedding(x_support_set_task, 0, training=training_phase) x_target_set_task = self.embedding(x_target_set_task, 0, training=training_phase) for num_step in range(num_steps): # operates on the support set support_loss, support_preds = self.net_forward( x=x_support_set_task, y=y_support_set_task, weights=names_weights_copy, backup_running_statistics=True if (num_step == 0) else False, training=True, num_step=num_step) # this is update of theta from the copy of current theta_0 and onward # i.e. inner loop optimization wrt support set names_weights_copy = self.apply_inner_loop_update( loss=support_loss, names_weights_copy=names_weights_copy, use_second_order=use_second_order, current_step_idx=num_step) # TODO: inner loop OPTIMIZATION wrt target set??? if use_multi_step_loss_optimization and training_phase and epoch < self.args.multi_step_loss_num_epochs: # this is MAML++ way target_loss, target_preds = self.net_forward( x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) task_losses.append( per_step_loss_importance_vectors[num_step] * target_loss) else: if num_step == ( self.args.number_of_training_steps_per_iter - 1): target_loss, target_preds = self.net_forward( x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) task_losses.append(target_loss) if self.args.use_critic: for i in range(self.args.num_critic_updates): # TODO: here must be an update using the Critic (start without g) # F = {f(x^b_T, θ_{N+j}), θ_{N+j}, g(xS, xn)} # θ_{N+j+1} = θ_{N+j} − \gamma * \nabla_{θ_{N+j}} C(F,W) critic_loss, target_preds = self.net_forward_critic( x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step + i) names_weights_copy = self.apply_inner_loop_update( loss=critic_loss, names_weights_copy=names_weights_copy, use_second_order=use_second_order, current_step_idx=num_step + i) target_loss, target_preds = self.net_forward( x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) task_losses.append(target_loss) per_task_target_preds[task_id] = target_preds.detach().cpu().numpy( ) _, predicted = torch.max(target_preds.data, 1) accuracy = predicted.float().eq( y_target_set_task.data.float()).cpu().float() task_losses = torch.sum(torch.stack(task_losses)) total_losses.append(task_losses) total_accuracies.extend(accuracy) if not training_phase: self.classifier.restore_backup_stats() losses = self.get_across_task_loss_metrics( total_losses=total_losses, total_accuracies=total_accuracies) for idx, item in enumerate(per_step_loss_importance_vectors): losses['loss_importance_vector_{}'.format( idx)] = item.detach().cpu().numpy() return losses, per_task_target_preds def net_forward(self, x, y, weights, backup_running_statistics, training, num_step): """ A base model forward pass on some data points x. Using the parameters in the weights dictionary. Also requires boolean flags indicating whether to reset the running statistics at the end of the run (if at evaluation phase). A flag indicating whether this is the training session and an int indicating the current step's number in the inner loop. :param x: A data batch of shape b, c, h, w :param y: A data targets batch of shape b, n_classes :param weights: A dictionary containing the weights to pass to the network. :param backup_running_statistics: A flag indicating whether to reset the batch norm running statistics to their previous values after the run (only for evaluation) :param training: A flag indicating whether the current process phase is a training or evaluation. :param num_step: An integer indicating the number of the step in the inner loop. :return: the crossentropy losses with respect to the given y, the predictions of the base model. """ preds = self.classifier.forward( x=x, params=weights, training=training, backup_running_statistics=backup_running_statistics, num_step=num_step) loss = F.cross_entropy(input=preds, target=y) return loss, preds def net_forward_critic(self, x, y, weights, backup_running_statistics, training, num_step): """ A base model forward pass on some data points x. Using the parameters in the weights dictionary. Also requires boolean flags indicating whether to reset the running statistics at the end of the run (if at evaluation phase). A flag indicating whether this is the training session and an int indicating the current step's number in the inner loop. :param x: A data batch of shape b, c, h, w :param y: A data targets batch of shape b, n_classes :param weights: A dictionary containing the weights to pass to the network. :param backup_running_statistics: A flag indicating whether to reset the batch norm running statistics to their previous values after the run (only for evaluation) :param training: A flag indicating whether the current process phase is a training or evaluation. :param num_step: An integer indicating the number of the step in the inner loop. :return: the crossentropy losses with respect to the given y, the predictions of the base model. """ preds = self.classifier.forward( x=x, params=weights, training=training, backup_running_statistics=backup_running_statistics, num_step=num_step) print(weights.keys()) params1d = torch.cat( [torch.reshape(p, (1, -1)) for p in list(weights.values())], dim=1) print(params1d.shape) loss = self.critic(preds, params1d) return loss, preds def trainable_parameters(self): """ Returns an iterator over the trainable parameters of the model. """ for param in self.parameters(): if param.requires_grad: yield param def train_forward_prop(self, data_batch, epoch): """ Runs an outer loop forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, per_task_target_preds = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=self.args.second_order and epoch > self.args.first_order_to_second_order_epoch, use_multi_step_loss_optimization=self.args. use_multi_step_loss_optimization, num_steps=self.args.number_of_training_steps_per_iter, training_phase=True) return losses, per_task_target_preds def evaluation_forward_prop(self, data_batch, epoch): """ Runs an outer loop evaluation forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, per_task_target_preds = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=False, use_multi_step_loss_optimization=True, num_steps=self.args.number_of_evaluation_steps_per_iter, training_phase=False) return losses, per_task_target_preds def meta_update(self, loss): """ Applies an outer loop update on the meta-parameters of the model. :param loss: The current crossentropy loss. """ self.optimizer.zero_grad() loss.backward(retain_graph=True) if 'imagenet' in self.args.dataset_name: for name, param in self.classifier.named_parameters(): if param.requires_grad: param.grad.data.clamp_( -10, 10 ) # not sure if this is necessary, more experiments are needed if self.args.high_end: for name, param in self.embedding.named_parameters(): if param.requires_grad: param.grad.data.clamp_( -10, 10 ) # not sure if this is necessary, more experiments are needed self.optimizer.step() def critic_meta_update(self, loss): self.critic_optimizer.zero_grad() loss.backward() self.critic_optimizer.step() def run_train_iter(self, data_batch, epoch): """ Runs an outer loop update step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ epoch = int(epoch) self.scheduler.step(epoch=epoch) if self.current_epoch != epoch: self.current_epoch = epoch if not self.training: self.train() x_support_set, x_target_set, y_support_set, y_target_set = data_batch x_support_set = torch.Tensor(x_support_set).float().to( device=self.device) x_target_set = torch.Tensor(x_target_set).float().to( device=self.device) y_support_set = torch.Tensor(y_support_set).long().to( device=self.device) y_target_set = torch.Tensor(y_target_set).long().to(device=self.device) data_batch = (x_support_set, x_target_set, y_support_set, y_target_set) losses, per_task_target_preds = self.train_forward_prop( data_batch=data_batch, epoch=epoch) self.meta_update(loss=losses['loss']) if self.args.use_critic: self.critic_meta_update(loss=losses['loss']) losses['learning_rate'] = self.scheduler.get_lr()[0] self.optimizer.zero_grad() self.zero_grad() return losses, per_task_target_preds def run_validation_iter(self, data_batch): """ Runs an outer loop evaluation step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ if self.training: self.eval() x_support_set, x_target_set, y_support_set, y_target_set = data_batch x_support_set = torch.Tensor(x_support_set).float().to( device=self.device) x_target_set = torch.Tensor(x_target_set).float().to( device=self.device) y_support_set = torch.Tensor(y_support_set).long().to( device=self.device) y_target_set = torch.Tensor(y_target_set).long().to(device=self.device) data_batch = (x_support_set, x_target_set, y_support_set, y_target_set) losses, per_task_target_preds = self.evaluation_forward_prop( data_batch=data_batch, epoch=self.current_epoch) # losses['loss'].backward() # uncomment if you get the weird memory error # self.zero_grad() # self.optimizer.zero_grad() return losses, per_task_target_preds def save_model(self, model_save_dir, state): """ Save the network parameter state and experiment state dictionary. :param model_save_dir: The directory to store the state at. :param state: The state containing the experiment state and the network. It's in the form of a dictionary object. """ state['network'] = self.state_dict() torch.save(state, f=model_save_dir) def load_model(self, model_save_dir, model_name, model_idx): """ Load checkpoint and return the state dictionary containing the network state params and experiment state. :param model_save_dir: The directory from which to load the files. :param model_name: The model_name to be loaded from the direcotry. :param model_idx: The index of the model (i.e. epoch number or 'latest' for the latest saved model of the current experiment) :return: A dictionary containing the experiment state and the saved model parameters. """ filepath = os.path.join(model_save_dir, "{}_{}".format(model_name, model_idx)) state = torch.load(filepath) state_dict_loaded = state['network'] self.load_state_dict(state_dict=state_dict_loaded) return state
def __init__(self, im_shape, device, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLFewShotClassifier, self).__init__() self.args = args self.device = device self.batch_size = args.batch_size self.use_cuda = args.use_cuda self.im_shape = im_shape self.current_epoch = 0 self.rng = set_torch_seed(seed=args.seed) if args.high_end: self.embedding = HighEndEmbedding(device, args, 3).to(device) self.classifier = HighEndClassifier( device, args, self.embedding.n_out_channels).to(device) else: self.classifier = VGGReLUNormNetwork( im_shape=self.im_shape, num_output_classes=self.args.num_classes_per_set, args=args, device=device, meta_classifier=True).to(device=self.device) self.task_learning_rate = args.task_learning_rate self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=device, init_learning_rate=self.task_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter + self.args.num_critic_updates, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters())) print("Inner Loop parameters") for key, value in self.inner_loop_optimizer.named_parameters(): print(key, value.shape) if args.use_critic: print( sum([ reduce(mul, p.size(), 1) for p in list( self.get_inner_loop_parameter_dict( self.classifier.named_parameters()).values()) ])) self.critic = Critic(n_theta=sum([ reduce(mul, p.size(), 1) for p in list( self.get_inner_loop_parameter_dict( self.classifier.named_parameters()).values()) ])) self.use_cuda = args.use_cuda self.device = device self.args = args self.to(device) print("Outer Loop parameters") for name, param in self.named_parameters(): if param.requires_grad: print(name, param.shape, param.device, param.requires_grad) if args.high_end: self.optimizer = optim.SGD(self.trainable_parameters(), lr=1e-4) else: self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.meta_learning_rate, amsgrad=False) if args.use_critic: self.critic_optimizer = optim.SGD(self.critic.parameters(), lr=1e-6) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate)
class MAMLRegressor(nn.Module): def __init__(self, input_shape, device, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLRegressor, self).__init__() self.args = args self.device = device self.batch_size = args.batch_size self.use_cuda = args.use_cuda self.current_epoch = 0 self.input_shape = input_shape self.rng = set_torch_seed(seed=args.seed) self.args.rng = self.rng self.regressor = IFCNReLUNormNetwork(input_shape=self.input_shape, args=self.args, device=device, meta=True).to(device=self.device) self.task_learning_rate = args.task_learning_rate self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=device, init_learning_rate=self.task_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.regressor.named_parameters())) print("Inner Loop parameters") for key, value in self.inner_loop_optimizer.named_parameters(): print(key, value.shape) self.use_cuda = args.use_cuda self.device = device self.args = args self.to(device) print("Outer Loop parameters") for name, param in self.named_parameters(): if param.requires_grad: print(name, param.shape, param.device, param.requires_grad) self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.meta_learning_rate, amsgrad=False) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate) self.optimizer_adapt = optim.Adam( self.trainable_adaptation_parameters(), lr=args.meta_learning_rate, amsgrad=False) self.scheduler_adapt = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer_adapt, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate) def get_per_step_loss_importance_vector(self): """ Generates a tensor of dimensionality (num_inner_loop_steps) indicating the importance of each step's target loss towards the optimization loss. :return: A tensor to be used to compute the weighted average of the loss, useful for the MSL (Multi Step Loss) mechanism. """ loss_weights = np.ones( shape=(self.args.number_of_training_steps_per_iter)) * ( 1.0 / self.args.number_of_training_steps_per_iter) decay_rate = 1.0 / self.args.number_of_training_steps_per_iter / self.args.multi_step_loss_num_epochs min_value_for_non_final_losses = 0.03 / self.args.number_of_training_steps_per_iter for i in range(len(loss_weights) - 1): curr_value = np.maximum( loss_weights[i] - (self.current_epoch * decay_rate), min_value_for_non_final_losses) loss_weights[i] = curr_value curr_value = np.minimum( loss_weights[-1] + (self.current_epoch * (self.args.number_of_training_steps_per_iter - 1) * decay_rate), 1.0 - ((self.args.number_of_training_steps_per_iter - 1) * min_value_for_non_final_losses)) loss_weights[-1] = curr_value loss_weights = torch.Tensor(loss_weights).to(device=self.device) return loss_weights def get_inner_loop_parameter_dict(self, params): """ Returns a dictionary with the parameters to use for inner loop updates. :param params: A dictionary of the network's parameters. :return: A dictionary of the parameters to use for the inner loop optimization process. """ param_dict = dict() for name, param in params: if param.requires_grad: if self.args.enable_inner_loop_optimizable_bn_params: param_dict[name] = param.to(device=self.device) else: # if "norm_layer" not in name: # param_dict[name] = param.to(device=self.device) if "layer_dict.linear.bias" in name or "layer_dict.linear.weights" in name: param_dict[name] = param.to(device=self.device) return param_dict def apply_inner_loop_update(self, loss, names_weights_copy, use_second_order, current_step_idx): """ Applies an inner loop update given current step's loss, the weights to update, a flag indicating whether to use second order derivatives and the current step's index. :param loss: Current step's loss with respect to the support set. :param names_weights_copy: A dictionary with names to parameters to update. :param use_second_order: A boolean flag of whether to use second order derivatives. :param current_step_idx: Current step's index. :return: A dictionary with the updated weights (name, param) """ self.regressor.zero_grad(names_weights_copy) grads = torch.autograd.grad(loss, names_weights_copy.values(), create_graph=use_second_order) names_grads_wrt_params = dict(zip(names_weights_copy.keys(), grads)) names_weights_copy = self.inner_loop_optimizer.update_params( names_weights_dict=names_weights_copy, names_grads_wrt_params_dict=names_grads_wrt_params, num_step=current_step_idx) return names_weights_copy def get_across_task_loss_metrics(self, total_losses, total_accuracies): losses = dict() losses['loss'] = torch.mean(torch.stack(total_losses[0])) losses['loss_supp'] = torch.mean(torch.stack(total_losses[1])) losses['loss_0'] = torch.mean(torch.stack(total_losses[2])) losses['loss_supp_0'] = torch.mean(torch.stack(total_losses[3])) if sum(~np.isnan(total_accuracies[0])) == 0: losses['accuracy'] = 0 else: losses['accuracy'] = np.mean( np.array(total_accuracies[0])[~np.isnan(total_accuracies[0])]) losses['accuracy_supp'] = np.mean( np.array(total_accuracies[1])[~np.isnan(total_accuracies[1])]) losses['accuracy_0'] = np.mean( np.array(total_accuracies[2])[~np.isnan(total_accuracies[2])]) losses['accuracy_supp_0'] = np.mean( np.array(total_accuracies[3])[~np.isnan(total_accuracies[3])]) return losses def forward(self, data_batch, epoch, use_second_order, use_multi_step_loss_optimization, num_steps, training_phase): """ Runs a forward outer loop pass on the batch of tasks using the MAML/++ framework. :param data_batch: A data batch containing the support and target sets. :param epoch: Current epoch's index :param use_second_order: A boolean saying whether to use second order derivatives. :param use_multi_step_loss_optimization: Whether to optimize on the outer loop using just the last step's target loss (True) or whether to use multi step loss which improves the stability of the system (False) :param num_steps: Number of inner loop steps. :param training_phase: Whether this is a training phase (True) or an evaluation phase (False) :return: A dictionary with the collected losses of the current outer forward propagation. """ support_set_x, support_set_y, support_set_z, support_set_assay, \ target_set_x, target_set_y, target_set_z, target_set_assay = data_batch b = len(support_set_y) total_losses = [] total_accuracies = [] total_losses_supp = [] total_accuracies_supp = [] total_losses_0 = [] total_accuracies_0 = [] total_losses_supp_0 = [] total_accuracies_supp_0 = [] total_losses_only = [] per_task_target_preds = [[] for i in range(len(target_set_x))] self.regressor.zero_grad() for task_id, (support_set_x_task, support_set_y_task, support_set_z_task, support_set_assay_task, target_set_x_task, target_set_y_task, target_set_z_task, target_set_assay_task,) in \ enumerate(zip(support_set_x, support_set_y, support_set_z, support_set_assay, target_set_x, target_set_y, target_set_z, target_set_assay)): # first of all, put all tensors to the device support_set_x_task = torch.Tensor( support_set_x_task[0]).float().to(device=self.device) support_set_y_task = torch.Tensor( support_set_y_task[0]).float().to(device=self.device) support_set_z_task = torch.IntTensor( support_set_z_task[0]).int().to(device=self.device) support_set_assay_task = torch.LongTensor( support_set_assay_task).int().to(device=self.device) target_set_x_task = torch.Tensor( target_set_x_task[0]).float().to(device=self.device) target_set_y_task = torch.Tensor( target_set_y_task[0]).float().to(device=self.device) target_set_z_task = torch.IntTensor( target_set_z_task[0]).int().to(device=self.device) target_set_assay_task = torch.LongTensor( target_set_assay_task).int().to(device=self.device) task_losses = [] task_accuracies = [] task_losses_supp = [] task_accuracies_supp = [] task_losses_0 = [] task_accuracies_0 = [] task_losses_supp_0 = [] task_accuracies_supp_0 = [] task_losses_only = [] per_step_loss_importance_vectors = self.get_per_step_loss_importance_vector( ) names_weights_copy = self.get_inner_loop_parameter_dict( self.regressor.named_parameters()) ns, fp_dim = support_set_x_task.shape nt, _ = target_set_x_task.shape r = np.random.uniform(-3, 3) for num_step in range(num_steps): support_loss, support_preds = self.net_forward( x=support_set_x_task, y=support_set_y_task + r, weights=names_weights_copy, backup_running_statistics=True if (num_step == 0) else False, training=True, num_step=num_step) # , support=True) if num_step == 0: target_loss, target_preds = self.net_forward( x=target_set_x_task, y=target_set_y_task + r, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=0) task_losses_supp_0.append(support_loss) task_losses_0.append(target_loss) accuracy_0 = pearson_score( target_set_y_task.detach().cpu().numpy(), (target_preds - r).detach().cpu().numpy()) task_losses_0 = torch.sum(torch.stack(task_losses_0)) total_losses_0.append(task_losses_0) total_accuracies_0.append(accuracy_0) accuracy_supp_0 = pearson_score( support_set_y_task.detach().cpu().numpy(), (support_preds - r).detach().cpu().numpy()) task_losses_supp_0 = torch.sum( torch.stack(task_losses_supp_0)) total_losses_supp_0.append(task_losses_supp_0) total_accuracies_supp_0.append(accuracy_supp_0) names_weights_copy = self.apply_inner_loop_update( loss=support_loss, names_weights_copy=names_weights_copy, use_second_order=use_second_order, current_step_idx=num_step) support_loss, support_preds = self.net_forward( x=support_set_x_task, y=support_set_y_task + r, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_steps - 1) target_loss, target_preds = self.net_forward( x=target_set_x_task, y=target_set_y_task + r, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_steps - 1) task_losses.append(target_loss) task_losses_supp.append(support_loss) per_task_target_preds[task_id] = target_preds.detach().cpu().numpy( ) # _, predicted = torch.max(target_preds.data, 1) accuracy = pearson_score(target_set_y_task.detach().cpu().numpy(), (target_preds - r).detach().cpu().numpy()) task_losses = torch.sum(torch.stack(task_losses)) total_losses.append(task_losses) total_accuracies.append(accuracy) supp_accuracy = pearson_score( support_set_y_task.detach().cpu().numpy(), (support_preds - r).detach().cpu().numpy()) task_losses_supp = torch.sum(torch.stack(task_losses_supp)) total_losses_supp.append(task_losses_supp) total_accuracies_supp.append(supp_accuracy) # pdb.set_trace() if not training_phase: self.regressor.restore_backup_stats() losses = self.get_across_task_loss_metrics(total_losses=[ total_losses, total_losses_supp, total_losses_0, total_losses_supp_0 ], total_accuracies=[ total_accuracies, total_accuracies_supp, total_accuracies_0, total_accuracies_supp_0 ]) for idx, item in enumerate(per_step_loss_importance_vectors): losses['loss_importance_vector_{}'.format( idx)] = item.detach().cpu().numpy() return losses, per_task_target_preds def net_forward(self, x, y, weights, backup_running_statistics, training, num_step, mixup=None, lam=None, support=None): """ A base model forward pass on some data points x. Using the parameters in the weights dictionary. Also requires boolean flags indicating whether to reset the running statistics at the end of the run (if at evaluation phase). A flag indicating whether this is the training session and an int indicating the current step's number in the inner loop. :param x: A data batch of shape b, c, h, w :param y: A data targets batch of shape b, n_classes :param weights: A dictionary containing the weights to pass to the network. :param backup_running_statistics: A flag indicating whether to reset the batch norm running statistics to their previous values after the run (only for evaluation) :param training: A flag indicating whether the current process phase is a training or evaluation. :param num_step: An integer indicating the number of the step in the inner loop. :return: the crossentropy losses with respect to the given y, the predictions of the base model. """ preds = self.regressor.forward( x=x, params=weights, training=training, backup_running_statistics=backup_running_statistics, num_step=num_step, mixup=mixup, lam=lam) if mixup is not None: npreds = preds.shape[0] preds = preds[:int(npreds / 2), :] # loss = F.cross_entropy(input=preds, target=y) # if support is not None: # loss = F.mse_loss(input=preds * sample_weights, target=y.unsqueeze(dim=-1)) # else: loss = F.mse_loss(input=preds, target=y.unsqueeze(dim=-1)) return loss, preds def trainable_parameters(self): """ Returns an iterator over the trainable parameters of the model. """ for param in self.parameters(): if param.requires_grad: yield param def trainable_feature_parameters(self): for name, param in self.named_parameters(): if param.requires_grad and 'learning_rates' not in name: # 'layer_dict.linear.bias' not in name and 'layer_dict.linear.weights' not in name and 'learning_rates' not in name: yield param def trainable_adaptation_parameters(self): for name, param in self.named_parameters(): if param.requires_grad and ('layer_dict.linear.bias' in name or 'layer_dict.linear.weights' in name or 'learning_rates' in name): yield param def train_forward_prop(self, data_batch, epoch): """ Runs an outer loop forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, per_task_target_preds = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=self.args.second_order and epoch > self.args.first_order_to_second_order_epoch, use_multi_step_loss_optimization=self.args. use_multi_step_loss_optimization, num_steps=self.args.number_of_training_steps_per_iter, training_phase=True) return losses, per_task_target_preds def evaluation_forward_prop(self, data_batch, epoch): """ Runs an outer loop evaluation forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, per_task_target_preds = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=False, use_multi_step_loss_optimization=True, num_steps=self.args.number_of_evaluation_steps_per_iter, training_phase=False) return losses, per_task_target_preds def meta_update(self, loss): """ Applies an outer loop update on the meta-parameters of the model. :param loss: The current crossentropy loss. """ self.optimizer.zero_grad() loss.backward() if 'imagenet' in self.args.dataset_name: for name, param in self.regressor.named_parameters(): if param.requires_grad: param.grad.data.clamp_( -10, 10 ) # not sure if this is necessary, more experiments are needed self.optimizer.step() def run_train_iter(self, data_batch, epoch): """ Runs an outer loop update step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ epoch = int(epoch) self.scheduler.step(epoch=epoch) if self.current_epoch != epoch: self.current_epoch = epoch if not self.training: self.train() losses, per_task_target_preds = self.train_forward_prop( data_batch=data_batch, epoch=epoch) self.meta_update(loss=losses['loss']) ''' if epoch < 10: self.optimizer.zero_grad() losses['loss'].backward() self.optimizer.step() else: self.optimizer_adapt.zero_grad() losses['loss'].backward() self.optimizer_adapt.step() ''' ''' self.optimizer.zero_grad() losses['loss'].backward(retain_graph=True) grad_loss_backup = {} for name, param in self.named_parameters(): if param.requires_grad and name not in grad_loss_backup and 'learning_rates' not in name and "layer_dict.linear.bias" not in name or "layer_dict.linear.weights" not in name: grad_loss_backup[name] = copy.deepcopy(param.grad) self.optimizer.zero_grad() losses['loss_only'].backward() for name, param in self.named_parameters(): if param.requires_grad and 'learning_rates' not in name and "layer_dict.linear.bias" not in name or "layer_dict.linear.weights" not in name: param.grad = copy.deepcopy(grad_loss_backup[name]) self.optimizer.step() ''' losses['learning_rate'] = self.scheduler.get_lr()[0] self.optimizer.zero_grad() self.optimizer_adapt.zero_grad() self.zero_grad() return losses, per_task_target_preds def run_validation_iter(self, data_batch): """ Runs an outer loop evaluation step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ if self.training: self.eval() losses, per_task_target_preds = self.evaluation_forward_prop( data_batch=data_batch, epoch=self.current_epoch) # losses['loss'].backward() # uncomment if you get the weird memory error # self.zero_grad() # self.optimizer.zero_grad() return losses, per_task_target_preds def save_model(self, model_save_dir, state): """ Save the network parameter state and experiment state dictionary. :param model_save_dir: The directory to store the state at. :param state: The state containing the experiment state and the network. It's in the form of a dictionary object. """ state['network'] = self.state_dict() torch.save(state, f=model_save_dir) def load_model(self, model_save_dir, model_name, model_idx): """ Load checkpoint and return the state dictionary containing the network state params and experiment state. :param model_save_dir: The directory from which to load the files. :param model_name: The model_name to be loaded from the direcotry. :param model_idx: The index of the model (i.e. epoch number or 'latest' for the latest saved model of the current experiment) :return: A dictionary containing the experiment state and the saved model parameters. """ filepath = os.path.join(model_save_dir, "{}_{}".format(model_name, model_idx)) state = torch.load(filepath) state_dict_loaded = state['network'] self.load_state_dict(state_dict=state_dict_loaded) return state
def __init__(self, input_shape, device, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLRegressor, self).__init__() self.args = args self.device = device self.batch_size = args.batch_size self.use_cuda = args.use_cuda self.current_epoch = 0 self.input_shape = input_shape self.rng = set_torch_seed(seed=args.seed) self.args.rng = self.rng self.regressor = IFCNReLUNormNetwork(input_shape=self.input_shape, args=self.args, device=device, meta=True).to(device=self.device) self.task_learning_rate = args.task_learning_rate self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=device, init_learning_rate=self.task_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.regressor.named_parameters())) print("Inner Loop parameters") for key, value in self.inner_loop_optimizer.named_parameters(): print(key, value.shape) self.use_cuda = args.use_cuda self.device = device self.args = args self.to(device) print("Outer Loop parameters") for name, param in self.named_parameters(): if param.requires_grad: print(name, param.shape, param.device, param.requires_grad) self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.meta_learning_rate, amsgrad=False) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate) self.optimizer_adapt = optim.Adam( self.trainable_adaptation_parameters(), lr=args.meta_learning_rate, amsgrad=False) self.scheduler_adapt = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer_adapt, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate)
def __init__(self, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(SceneAdaptiveInterpolation, self).__init__() self.args = args self.device = torch.device('cuda') if args.cuda else torch.device( 'cpu') self.batch_size = args.batch_size self.use_cuda = args.cuda # self.im_shape = im_shape self.current_epoch = 0 self.rng = set_torch_seed(seed=args.random_seed) # self.classifier = VGGReLUNormNetwork(im_shape=self.im_shape, num_output_classes=self.args.num_classes_per_set, # args=args, device=device, meta_classifier=True).to(device=self.device) if self.args.model == 'sepconv': print('Building SepConv model...') from sepconv.model import MetaNetwork as MetaSepConv self.net = MetaSepConv(resume=False if self.args.resume else True, strModel='l1').to(self.device) elif self.args.model == 'cain': print('Building CAIN model...') from cain.model import MetaCAIN self.net = MetaCAIN(depth=3, resume=False if self.args.resume else True).to( self.device) elif self.args.model == 'superslomo': print('Building SuperSloMo model...') from superslomo.model import MetaSuperSloMo self.net = MetaSuperSloMo( self.device, resume=False if self.args.resume else True).to(self.device) # reverse normalization to transform super-slomo outputs to 0~1 scale neg_mean = [-.429, -0.431, -0.397] std = [1, 1, 1] self.revNormalize = transforms.Normalize(mean=neg_mean, std=std) elif args.model == 'voxelflow': print('Building Deep VoxelFlow (DVF) model...') from voxelflow.core.models.voxel_flow import MetaVoxelFlow self.net = MetaVoxelFlow( self.args, resume=False if self.args.resume else True).to(self.device) else: raise NotImplementedError('Model not implemented yet!') self.inner_learning_rate = args.inner_lr if self.args.metasgd: print('Adaptation with Meta-SGD') self.inner_loop_optimizer = MetaSGDLearningRule( device=self.device, optimizer=self.args.optimizer, init_learning_rate=self.inner_learning_rate) else: self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=self.device, optimizer=self.args.optimizer, #'Adamax', init_learning_rate=self.inner_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) names_weights_dict = self.get_inner_loop_parameter_dict( params=self.net.named_parameters()) self.inner_loop_optimizer.initialize( names_weights_dict=names_weights_dict) # Attenuator for L2F if self.args.attenuate: num_layers = len(names_weights_dict.keys()) print('# of layers: %d' % num_layers) self.attenuator = nn.Sequential( nn.Linear(num_layers, num_layers), nn.ReLU(inplace=True), nn.Linear(num_layers, num_layers), nn.Sigmoid()).to(device=self.device) # initialize to output zero self.gamma_mult = nn.Parameter(torch.zeros(1)) # print("Inner Loop parameters") # for key, value in self.inner_loop_optimizer.named_parameters(): # print(key, value.shape) self.use_cuda = args.cuda self.args = args self.to(self.device) # print("Outer Loop parameters") # for name, param in self.named_parameters(): # if param.requires_grad: # print(name, param.shape, param.device, param.requires_grad) if self.args.optimizer == 'Adam': print('Using optimizer Adam.') if self.args.model == 'voxelflow': policies = self.net.get_optim_policies() self.optimizer = optim.Adam( policies, lr=args.outer_lr, weight_decay=args.weight_decay) #Optim(policies, args) else: self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.outer_lr, betas=(0.9, 0.99)) elif self.args.optimizer == 'Adamax': print('Using optimizer Adamax.') self.optimizer = optim.Adamax(self.trainable_parameters(), lr=args.outer_lr, betas=(0.9, 0.999)) else: self.optimizer = optim.SGD(self.trainable_parameters, lr=args.outer_lr) self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer=self.optimizer, mode='min', factor=0.2, patience=5, verbose=True) num_params = 0 for param in list(self.trainable_parameters()): # print(param.shape) num_params += param.numel() print('# of parameters: %d' % num_params) self.criterion = Loss(args) if self.args.resume: print('Resume training') utils.load_checkpoint(args, self, None) if self.args.pretrained_model is not None: print('Loading pretrained model: %s' % self.args.pretrained_model) checkpoint = torch.load(self.args.pretrained_model) if self.args.model == 'superslomo': if 'state_dictFC' in checkpoint.keys(): utils.lossy_load_state_dict(self.net.flowComp, checkpoint['state_dictFC']) utils.lossy_load_state_dict(self.net.arbTimeFlowIntrp, checkpoint['state_dictAT']) else: utils.lossy_load_state_dict(self, checkpoint['state_dict']) else: utils.lossy_load_state_dict(self.net, checkpoint['state_dict'])
class SceneAdaptiveInterpolation(nn.Module): # def __init__(self, im_shape, device, args): def __init__(self, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(SceneAdaptiveInterpolation, self).__init__() self.args = args self.device = torch.device('cuda') if args.cuda else torch.device( 'cpu') self.batch_size = args.batch_size self.use_cuda = args.cuda # self.im_shape = im_shape self.current_epoch = 0 self.rng = set_torch_seed(seed=args.random_seed) # self.classifier = VGGReLUNormNetwork(im_shape=self.im_shape, num_output_classes=self.args.num_classes_per_set, # args=args, device=device, meta_classifier=True).to(device=self.device) if self.args.model == 'sepconv': print('Building SepConv model...') from sepconv.model import MetaNetwork as MetaSepConv self.net = MetaSepConv(resume=False if self.args.resume else True, strModel='l1').to(self.device) elif self.args.model == 'cain': print('Building CAIN model...') from cain.model import MetaCAIN self.net = MetaCAIN(depth=3, resume=False if self.args.resume else True).to( self.device) elif self.args.model == 'superslomo': print('Building SuperSloMo model...') from superslomo.model import MetaSuperSloMo self.net = MetaSuperSloMo( self.device, resume=False if self.args.resume else True).to(self.device) # reverse normalization to transform super-slomo outputs to 0~1 scale neg_mean = [-.429, -0.431, -0.397] std = [1, 1, 1] self.revNormalize = transforms.Normalize(mean=neg_mean, std=std) elif args.model == 'voxelflow': print('Building Deep VoxelFlow (DVF) model...') from voxelflow.core.models.voxel_flow import MetaVoxelFlow self.net = MetaVoxelFlow( self.args, resume=False if self.args.resume else True).to(self.device) else: raise NotImplementedError('Model not implemented yet!') self.inner_learning_rate = args.inner_lr if self.args.metasgd: print('Adaptation with Meta-SGD') self.inner_loop_optimizer = MetaSGDLearningRule( device=self.device, optimizer=self.args.optimizer, init_learning_rate=self.inner_learning_rate) else: self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=self.device, optimizer=self.args.optimizer, #'Adamax', init_learning_rate=self.inner_learning_rate, total_num_inner_loop_steps=self.args. number_of_training_steps_per_iter, use_learnable_learning_rates=self.args. learnable_per_layer_per_step_inner_loop_learning_rate) names_weights_dict = self.get_inner_loop_parameter_dict( params=self.net.named_parameters()) self.inner_loop_optimizer.initialize( names_weights_dict=names_weights_dict) # Attenuator for L2F if self.args.attenuate: num_layers = len(names_weights_dict.keys()) print('# of layers: %d' % num_layers) self.attenuator = nn.Sequential( nn.Linear(num_layers, num_layers), nn.ReLU(inplace=True), nn.Linear(num_layers, num_layers), nn.Sigmoid()).to(device=self.device) # initialize to output zero self.gamma_mult = nn.Parameter(torch.zeros(1)) # print("Inner Loop parameters") # for key, value in self.inner_loop_optimizer.named_parameters(): # print(key, value.shape) self.use_cuda = args.cuda self.args = args self.to(self.device) # print("Outer Loop parameters") # for name, param in self.named_parameters(): # if param.requires_grad: # print(name, param.shape, param.device, param.requires_grad) if self.args.optimizer == 'Adam': print('Using optimizer Adam.') if self.args.model == 'voxelflow': policies = self.net.get_optim_policies() self.optimizer = optim.Adam( policies, lr=args.outer_lr, weight_decay=args.weight_decay) #Optim(policies, args) else: self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.outer_lr, betas=(0.9, 0.99)) elif self.args.optimizer == 'Adamax': print('Using optimizer Adamax.') self.optimizer = optim.Adamax(self.trainable_parameters(), lr=args.outer_lr, betas=(0.9, 0.999)) else: self.optimizer = optim.SGD(self.trainable_parameters, lr=args.outer_lr) self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer=self.optimizer, mode='min', factor=0.2, patience=5, verbose=True) num_params = 0 for param in list(self.trainable_parameters()): # print(param.shape) num_params += param.numel() print('# of parameters: %d' % num_params) self.criterion = Loss(args) if self.args.resume: print('Resume training') utils.load_checkpoint(args, self, None) if self.args.pretrained_model is not None: print('Loading pretrained model: %s' % self.args.pretrained_model) checkpoint = torch.load(self.args.pretrained_model) if self.args.model == 'superslomo': if 'state_dictFC' in checkpoint.keys(): utils.lossy_load_state_dict(self.net.flowComp, checkpoint['state_dictFC']) utils.lossy_load_state_dict(self.net.arbTimeFlowIntrp, checkpoint['state_dictAT']) else: utils.lossy_load_state_dict(self, checkpoint['state_dict']) else: utils.lossy_load_state_dict(self.net, checkpoint['state_dict']) ###### Below script needed only for distributed training # self.device = torch.device('cpu') # if torch.cuda.is_available(): # if torch.cuda.device_count() > 1: # self.to(torch.cuda.current_device()) # self.classifier = nn.DataParallel(module=self.classifier) # else: # self.to(torch.cuda.current_device()) # self.device = torch.cuda.current_device() def get_per_step_loss_importance_vector(self): """ Generates a tensor of dimensionality (num_inner_loop_steps) indicating the importance of each step's target loss towards the optimization loss. :return: A tensor to be used to compute the weighted average of the loss, useful for the MSL (Multi Step Loss) mechanism. """ loss_weights = np.ones( shape=(self.args.number_of_training_steps_per_iter)) * ( 1.0 / self.args.number_of_training_steps_per_iter) decay_rate = 1.0 / self.args.number_of_training_steps_per_iter / self.args.multi_step_loss_num_epochs min_value_for_non_final_losses = 0.03 / self.args.number_of_training_steps_per_iter for i in range(len(loss_weights) - 1): curr_value = np.maximum( loss_weights[i] - (self.current_epoch * decay_rate), min_value_for_non_final_losses) loss_weights[i] = curr_value curr_value = np.minimum( loss_weights[-1] + (self.current_epoch * (self.args.number_of_training_steps_per_iter - 1) * decay_rate), 1.0 - ((self.args.number_of_training_steps_per_iter - 1) * min_value_for_non_final_losses)) loss_weights[-1] = curr_value loss_weights = torch.Tensor(loss_weights).to(device=self.device) return loss_weights def get_inner_loop_parameter_dict(self, params): """ Returns a dictionary with the parameters to use for inner loop updates. :param params: A dictionary of the network's parameters. :return: A dictionary of the parameters to use for the inner loop optimization process. """ param_dict = dict() for name, param in params: if param.requires_grad: if self.args.enable_inner_loop_optimizable_bn_params: param_dict[name] = param.to(device=self.device) else: if "norm_layer" not in name: param_dict[name] = param.to(device=self.device) return param_dict def get_task_embeddings(self, frames, task_id, names_weights_copy): if self.args.mode == 'test': support_idxs = [[0, 1, 2], [1, 2, 3]] else: support_idxs = [[0, 2, 4], [2, 4, 6] ] # frame indices: input[0, 2, 4, 6] --> output[3] # target_idx = [2, 3, 4] support_loss = 0 for ind in support_idxs: _loss, _ = self.net_forward( frame0=frames[ind[0]][task_id].unsqueeze(0), frame1=frames[ind[2]][task_id].unsqueeze(0), target=frames[ind[1]][task_id].unsqueeze(0), weights=names_weights_copy, backup_running_statistics=True, training=True, num_step=0) support_loss = support_loss + _loss['total'] self.net.zero_grad(names_weights_copy) grads = torch.autograd.grad(support_loss, names_weights_copy.values(), create_graph=False, allow_unused=True) layerwise_mean_grads = [] for i in range(len(grads)): layerwise_mean_grads.append(grads[i].mean()) layerwise_mean_grads = torch.stack(layerwise_mean_grads) return layerwise_mean_grads def attenuate_init(self, task_embeddings, names_weights_copy): #gamma = 0.5 + self.attenuator(task_embeddings) # 0.5 is added to initialize gamma to 1 gamma = 1 - self.gamma_mult * self.attenuator(task_embeddings) gamma.clamp_(0, 1) gammas = [] for i in range(gamma.size(0)): gammas.append(gamma[i]) #print(gamma[i].item()) updated_weights = list( map( lambda current_params, gamma: ((gamma) * current_params.to(device=self.device)), names_weights_copy.values(), gamma)) updated_names_weights_copy = dict( zip(names_weights_copy.keys(), updated_weights)) return updated_names_weights_copy def apply_inner_loop_update(self, loss, names_weights_copy, use_second_order, current_step_idx): """ Applies an inner loop update given current step's loss, the weights to update, a flag indicating whether to use second order derivatives and the current step's index. :param loss: Current step's loss with respect to the support set. :param names_weights_copy: A dictionary with names to parameters to update. :param use_second_order: A boolean flag of whether to use second order derivatives. :param current_step_idx: Current step's index. :return: A dictionary with the updated weights (name, param) """ num_gpus = torch.cuda.device_count() if num_gpus > 1: self.net.module.zero_grad(params=names_weights_copy) else: self.net.zero_grad(params=names_weights_copy) grads = torch.autograd.grad(loss, names_weights_copy.values(), create_graph=use_second_order, allow_unused=True) names_grads_copy = dict(zip(names_weights_copy.keys(), grads)) # names_weights_copy = {key: value[0] for key, value in names_weights_copy.items()} ###### This is needed for summing up gradients w.r.t. different GPUs when distributed training # for key, grad in names_grads_copy.items(): # if grad is None: # print('Grads not found for inner loop parameter', key) # names_grads_copy[key] = names_grads_copy[key].sum(dim=0) names_weights_copy = self.inner_loop_optimizer.update_params( names_weights_dict=names_weights_copy, names_grads_wrt_params_dict=names_grads_copy, num_step=current_step_idx) # loss.backward() # self.inner_loop_optimizer.step() # num_devices = torch.cuda.device_count() if torch.cuda.is_available() else 1 # names_weights_copy = { # name.replace('module.', ''): value.unsqueeze(0).repeat( # [num_devices] + [1 for i in range(len(value.shape))]) for # name, value in names_weights_copy.items()} # names_weights_copy = { # name.replace('module.', ''): value for name, value in names_weights_copy.items()} return names_weights_copy def update_loss_metrics(self, task_losses, target_loss): """ :param task_losses: accumulator dictionary to gather all losses for logging to TensorBoard (updated in-place) :param target_loss: current loss to be updated to task_losses """ for loss_key, loss_value in target_loss.items(): if loss_key not in task_losses.keys(): task_losses[loss_key] = utils.AverageMeter() task_losses[loss_key].update( loss_value.detach().cpu().data.numpy()) def get_across_task_loss_metrics(self, total_losses, specific_losses): losses = dict() losses['loss'] = torch.mean(torch.stack(total_losses)) # losses['accuracy'] = np.mean(total_accuracies) for key, avg_meters in specific_losses.items(): losses[key] = avg_meters.avg return losses def forward(self, data_batch, epoch, use_second_order, use_multi_step_loss_optimization, num_steps, training_phase, do_evaluation=False): """ Runs a forward outer loop pass on the batch of tasks using the MAML/++ framework. :param data_batch: A data batch containing the support and target sets. :param epoch: Current epoch's index :param use_second_order: A boolean saying whether to use second order derivatives. :param use_multi_step_loss_optimization: Whether to optimize on the outer loop using just the last step's target loss (True) or whether to use multi step loss which improves the stability of the system (False) :param num_steps: Number of inner loop steps. :param training_phase: Whether this is a training phase (True) or an evaluation phase (False) :return: A dictionary with the collected losses of the current outer forward propagation. """ frames = data_batch total_losses = [] loss_accumulator = {'total': utils.AverageMeter()} metrics = {'psnr': utils.AverageMeter(), 'ssim': utils.AverageMeter()} per_task_target_preds = [[] for i in range(len(frames[0]))] self.net.zero_grad() for task_id in range(len(frames[0])): # loop over batch dimension task_losses = [] per_step_loss_importance_vectors = self.get_per_step_loss_importance_vector( ) names_weights_copy = self.get_inner_loop_parameter_dict( self.net.named_parameters()) names_weights_copy = { name.replace('module.', ''): value for name, value in names_weights_copy.items() } # inner loop support_idxs = [[0, 2, 4], [2, 4, 6] ] # frame indices: input[0, 2, 4, 6] --> output[3] target_idx = [2, 3, 4] self.inner_loop_optimizer.initialize_state() # Attenuate the initialization for L2F if self.args.attenuate: task_embeddings = self.get_task_embeddings( frames, task_id, names_weights_copy) names_weights_copy = self.attenuate_init( task_embeddings=task_embeddings, names_weights_copy=names_weights_copy) for num_step in range(num_steps): support_loss = 0 for ind in support_idxs: _loss, _ = self.net_forward( frame0=frames[ind[0]][task_id].unsqueeze(0), frame1=frames[ind[2]][task_id].unsqueeze(0), target=frames[ind[1]][task_id].unsqueeze(0), weights=names_weights_copy, backup_running_statistics=True if (num_step == 0) else False, training=True, num_step=num_step) support_loss = support_loss + _loss['total'] names_weights_copy = self.apply_inner_loop_update( loss=support_loss, names_weights_copy=names_weights_copy, use_second_order=use_second_order, current_step_idx=num_step) kwargs = { 'backup_running_statistics': False, 'training': True, 'num_step': num_step } if use_multi_step_loss_optimization and training_phase and epoch < self.args.multi_step_loss_num_epochs: target_loss, target_preds = self.net_forward( frame0=frames[target_idx[0]][task_id].unsqueeze(0), frame1=frames[target_idx[2]][task_id].unsqueeze(0), target=frames[target_idx[1]][task_id].unsqueeze(0), weights=names_weights_copy, **kwargs) task_losses.append( per_step_loss_importance_vectors[num_step] * target_loss['total']) self.update_loss_metrics(loss_accumulator, target_loss) if not (use_multi_step_loss_optimization and training_phase and epoch < self.args.multi_step_loss_num_epochs): kwargs = { 'backup_running_statistics': False, 'training': True, 'num_step': num_steps } target_loss, target_preds = self.net_forward( frame0=frames[target_idx[0]][task_id].unsqueeze(0), frame1=frames[target_idx[2]][task_id].unsqueeze(0), target=frames[target_idx[1]][task_id].unsqueeze(0), weights=names_weights_copy, **kwargs) task_losses.append(target_loss['total']) self.update_loss_metrics(loss_accumulator, target_loss) if self.args.model == 'superslomo': per_task_target_preds[task_id] = self.revNormalize( target_preds.detach().squeeze(0)).unsqueeze(0) else: per_task_target_preds[task_id] = target_preds.detach( ) # target_preds.shape: (1, C, H, W) if do_evaluation: if self.args.model == 'superslomo': output = self.revNormalize(target_preds.squeeze(0)) target = self.revNormalize(frames[target_idx[1]][task_id]) else: output = target_preds.squeeze(0) target = frames[target_idx[1]][task_id] output = output.detach() target = target.detach() psnr, ssim = utils.calc_metrics(output, target) # print(psnr, ssim) metrics['psnr'].update(psnr) metrics['ssim'].update(ssim) else: pass task_losses = torch.sum(torch.stack(task_losses)) total_losses.append(task_losses) if not training_phase: self.net.restore_backup_stats() losses = self.get_across_task_loss_metrics( total_losses=total_losses, specific_losses=loss_accumulator) for idx, item in enumerate(per_step_loss_importance_vectors): losses['loss_importance_vector_{}'.format( idx)] = item.detach().cpu().numpy() return losses, per_task_target_preds, metrics def net_forward(self, frame0, frame1, target, weights, backup_running_statistics, training, num_step): """ A base model forward pass on the input frames. Using the parameters in the weights dictionary. Also requires boolean flags indicating whether to reset the running statistics at the end of the run (if at evaluation phase). A flag indicating whether this is the training session and an int indicating the current step's number in the inner loop. :param frame0: A data batch containing the first input frames :param frame1: A data batch containing the second input frames :param weights: A dictionary containing the weights to pass to the network. :param backup_running_statistics: A flag indicating whether to reset the batch norm running statistics to their previous values after the run (only for evaluation) :param training: A flag indicating whether the current process phase is a training or evaluation. :param num_step: An integer indicating the number of the step in the inner loop. :return: the crossentropy losses with respect to the given y, the predictions of the base model. """ kwargs = { 'backup_running_statistics': backup_running_statistics, 'num_step': num_step } output = self.net.forward(frame0, frame1, params=weights, **kwargs) # output = self.net.forward(frame0, frame1, params=None, **kwargs) if self.args.model == 'superslomo': # output becomes a tuple output[1]['I0'], output[1]['I1'] = frame0, frame1 losses = self.criterion(output[0], target, **output[1]) output = output[0] else: losses = self.criterion(output, target) return losses, output def trainable_parameters(self): """ Returns an iterator over the trainable parameters of the model. """ for param in self.parameters(): if param.requires_grad: yield param def train_forward_prop(self, data_batch, epoch, do_evaluation=False): """ Runs an outer loop forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, preds, metrics = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=self.args.second_order and epoch > self.args.first_order_to_second_order_epoch, use_multi_step_loss_optimization=self.args. use_multi_step_loss_optimization, num_steps=self.args.number_of_training_steps_per_iter, training_phase=True, do_evaluation=do_evaluation) return losses, preds, metrics def evaluation_forward_prop(self, data_batch, epoch): """ Runs an outer loop evaluation forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, preds, metrics = self.forward( data_batch=data_batch, epoch=epoch, use_second_order=False, use_multi_step_loss_optimization=True, num_steps=self.args.number_of_evaluation_steps_per_iter, training_phase=False, do_evaluation=True) return losses, preds, metrics def meta_update(self, loss): """ Applies an outer loop update on the meta-parameters of the model. :param loss: The current crossentropy loss. """ self.optimizer.zero_grad() loss.backward() if False: #'imagenet' in self.args.dataset_name: for _, param in self.net.named_parameters(): if param.requires_grad: param.grad.data.clamp_( -10, 10 ) # not sure if this is necessary, more experiments are needed #grads = torch.autograd.grad(loss, self.net.parameters()) #for j, param in enumerate(self.net.parameters()): # param.grad = grads[j] self.optimizer.step() def run_train_iter(self, data_batch, epoch, do_evaluation=False): """ Runs an outer loop update step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ epoch = int(epoch) if self.current_epoch != epoch: self.current_epoch = epoch if not self.training: self.train() data_batch = [frame.to(device=self.device) for frame in data_batch] losses, preds, metrics = self.train_forward_prop( data_batch=data_batch, epoch=epoch, do_evaluation=do_evaluation) self.meta_update(loss=losses['loss']) self.optimizer.zero_grad() self.zero_grad() return losses, preds, metrics def run_validation_iter(self, data_batch): """ Runs an outer loop evaluation step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ if self.training: self.eval() data_batch = [frame.to(device=self.device) for frame in data_batch] losses, preds, metrics = self.evaluation_forward_prop( data_batch=data_batch, epoch=self.current_epoch) # losses['loss'].backward() # uncomment if you get the weird memory error # self.zero_grad() # self.optimizer.zero_grad() return losses, preds, metrics def run_test_iter(self, data_batch): """ Runs an outer loop evaluation step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ if self.training: self.eval() frames = [frame.to(device=self.device) for frame in data_batch] preds = [[] for i in range(len(frames[0]))] self.net.zero_grad() for task_id in range(len(frames[0])): # loop over batch dimension names_weights_copy = self.get_inner_loop_parameter_dict( self.net.named_parameters()) names_weights_copy = { name.replace('module.', ''): value for name, value in names_weights_copy.items() } # inner loop support_idxs = [[0, 1, 2], [1, 2, 3]] target_idx = [1, 2] # only input self.inner_loop_optimizer.initialize_state() # Attenuate the initialization for L2F if self.args.attenuate: task_embeddings = self.get_task_embeddings( frames, task_id, names_weights_copy) names_weights_copy = self.attenuate_init( task_embeddings=task_embeddings, names_weights_copy=names_weights_copy) for num_step in range( self.args.number_of_evaluation_steps_per_iter): support_loss = 0 for ind in support_idxs: _loss, _ = self.net_forward( frame0=frames[ind[0]][task_id].unsqueeze(0), frame1=frames[ind[2]][task_id].unsqueeze(0), target=frames[ind[1]][task_id].unsqueeze(0), weights=names_weights_copy, backup_running_statistics=True if (num_step == 0) else False, training=True, num_step=num_step) support_loss = support_loss + _loss['total'] names_weights_copy = self.apply_inner_loop_update( loss=support_loss, names_weights_copy=names_weights_copy, use_second_order=self.args.second_order, current_step_idx=num_step) frame0 = frames[target_idx[0]][task_id].unsqueeze(0) frame1 = frames[target_idx[1]][task_id].unsqueeze(0) kwargs = { 'backup_running_statistics': False, 'training': True, 'num_step': num_step } output = self.net.forward(frame0, frame1, params=names_weights_copy, **kwargs) if self.args.model == 'superslomo': # output becomes a tuple output[1]['I0'], output[1]['I1'] = frame0, frame1 output = self.revNormalize(output[0].squeeze(0)) else: output = output.squeeze(0) # print(output.shape) preds[task_id] = output.detach() # output.shape: (C, H, W) self.net.restore_backup_stats() return preds
class MAMLFewShotClassifier(nn.Module): def __init__(self, im_shape, device, args): """ Initializes a MAML few shot learning system :param im_shape: The images input size, in batch, c, h, w shape :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLFewShotClassifier, self).__init__() self.args = args self.device = device self.batch_size = args.batch_size self.use_cuda = args.use_cuda self.im_shape = im_shape self.current_epoch = 0 self.rng = set_torch_seed(seed=args.seed) self.classifier = VGGReLUNormNetwork(im_shape=self.im_shape, num_output_classes=self.args. num_classes_per_set, args=args, device=device, meta_classifier=True).to(device=self.device) self.task_learning_rate = args.task_learning_rate self.inner_loop_optimizer = LSLRGradientDescentLearningRule(device=device, init_learning_rate=self.task_learning_rate, total_num_inner_loop_steps=self.args.number_of_training_steps_per_iter, use_learnable_learning_rates=self.args.learnable_per_layer_per_step_inner_loop_learning_rate) self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict(params=self.classifier.named_parameters())) print("Inner Loop parameters") for key, value in self.inner_loop_optimizer.named_parameters(): print(key, value.shape) self.use_cuda = args.use_cuda self.device = device self.args = args self.to(device) print("Outer Loop parameters") for name, param in self.named_parameters(): if param.requires_grad: print(name, param.shape, param.device, param.requires_grad) self.optimizer = optim.Adam(self.trainable_parameters(), lr=args.meta_learning_rate, amsgrad=False) self.scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer=self.optimizer, T_max=self.args.total_epochs, eta_min=self.args.min_learning_rate) self.device = torch.device('cpu') if torch.cuda.is_available(): if torch.cuda.device_count() > 1: self.to(torch.cuda.current_device()) self.classifier = nn.DataParallel(module=self.classifier) else: self.to(torch.cuda.current_device()) self.device = torch.cuda.current_device() def get_per_step_loss_importance_vector(self): """ Generates a tensor of dimensionality (num_inner_loop_steps) indicating the importance of each step's target loss towards the optimization loss. :return: A tensor to be used to compute the weighted average of the loss, useful for the MSL (Multi Step Loss) mechanism. """ loss_weights = np.ones(shape=(self.args.number_of_training_steps_per_iter)) * ( 1.0 / self.args.number_of_training_steps_per_iter) decay_rate = 1.0 / self.args.number_of_training_steps_per_iter / self.args.multi_step_loss_num_epochs min_value_for_non_final_losses = 0.03 / self.args.number_of_training_steps_per_iter for i in range(len(loss_weights) - 1): curr_value = np.maximum(loss_weights[i] - (self.current_epoch * decay_rate), min_value_for_non_final_losses) loss_weights[i] = curr_value curr_value = np.minimum( loss_weights[-1] + (self.current_epoch * (self.args.number_of_training_steps_per_iter - 1) * decay_rate), 1.0 - ((self.args.number_of_training_steps_per_iter - 1) * min_value_for_non_final_losses)) loss_weights[-1] = curr_value loss_weights = torch.Tensor(loss_weights).to(device=self.device) return loss_weights def get_inner_loop_parameter_dict(self, params): """ Returns a dictionary with the parameters to use for inner loop updates. :param params: A dictionary of the network's parameters. :return: A dictionary of the parameters to use for the inner loop optimization process. """ param_dict = dict() for name, param in params: if param.requires_grad: if self.args.enable_inner_loop_optimizable_bn_params: param_dict[name] = param.to(device=self.device) else: if "norm_layer" not in name: param_dict[name] = param.to(device=self.device) return param_dict def apply_inner_loop_update(self, loss, names_weights_copy, use_second_order, current_step_idx): """ Applies an inner loop update given current step's loss, the weights to update, a flag indicating whether to use second order derivatives and the current step's index. :param loss: Current step's loss with respect to the support set. :param names_weights_copy: A dictionary with names to parameters to update. :param use_second_order: A boolean flag of whether to use second order derivatives. :param current_step_idx: Current step's index. :return: A dictionary with the updated weights (name, param) """ num_gpus = torch.cuda.device_count() if num_gpus > 1: self.classifier.module.zero_grad(params=names_weights_copy) else: self.classifier.zero_grad(params=names_weights_copy) grads = torch.autograd.grad(loss, names_weights_copy.values(), create_graph=use_second_order, allow_unused=True) names_grads_copy = dict(zip(names_weights_copy.keys(), grads)) names_weights_copy = {key: value[0] for key, value in names_weights_copy.items()} for key, grad in names_grads_copy.items(): if grad is None: print('Grads not found for inner loop parameter', key) names_grads_copy[key] = names_grads_copy[key].sum(dim=0) names_weights_copy = self.inner_loop_optimizer.update_params(names_weights_dict=names_weights_copy, names_grads_wrt_params_dict=names_grads_copy, num_step=current_step_idx) num_devices = torch.cuda.device_count() if torch.cuda.is_available() else 1 names_weights_copy = { name.replace('module.', ''): value.unsqueeze(0).repeat( [num_devices] + [1 for i in range(len(value.shape))]) for name, value in names_weights_copy.items()} return names_weights_copy def get_across_task_loss_metrics(self, total_losses, total_accuracies, total_auc, total_f1): losses = dict() metrics = dict() losses['loss'] = torch.mean(torch.stack(total_losses)) losses['accuracy'] = np.mean(total_accuracies) metrics['auc'] = np.mean(total_auc) metrics['f1'] = np.mean(total_f1) return losses, metrics def forward(self, data_batch, epoch, use_second_order, use_multi_step_loss_optimization, num_steps, training_phase): """ Runs a forward outer loop pass on the batch of tasks using the MAML/++ framework. :param data_batch: A data batch containing the support and target sets. :param epoch: Current epoch's index :param use_second_order: A boolean saying whether to use second order derivatives. :param use_multi_step_loss_optimization: Whether to optimize on the outer loop using just the last step's target loss (True) or whether to use multi step loss which improves the stability of the system (False) :param num_steps: Number of inner loop steps. :param training_phase: Whether this is a training phase (True) or an evaluation phase (False) :return: A dictionary with the collected losses of the current outer forward propagation. """ x_support_set, x_target_set, y_support_set, y_target_set = data_batch [b, ncs, spc] = y_support_set.shape self.num_classes_per_set = ncs total_losses = [] total_accuracies = [] total_auc = [] total_f1 = [] per_task_target_preds = [[] for i in range(len(x_target_set))] self.classifier.zero_grad() for task_id, (x_support_set_task, y_support_set_task, x_target_set_task, y_target_set_task) in \ enumerate(zip(x_support_set, y_support_set, x_target_set, y_target_set)): task_losses = [] task_accuracies = [] task_roc = [] task_f1 = [] per_step_loss_importance_vectors = self.get_per_step_loss_importance_vector() names_weights_copy = self.get_inner_loop_parameter_dict(self.classifier.named_parameters()) num_devices = torch.cuda.device_count() if torch.cuda.is_available() else 1 names_weights_copy = { name.replace('module.', ''): value.unsqueeze(0).repeat( [num_devices] + [1 for i in range(len(value.shape))]) for name, value in names_weights_copy.items()} n, s, c, h, w = x_target_set_task.shape x_support_set_task = x_support_set_task.view(-1, c, h, w) y_support_set_task = y_support_set_task.view(-1) x_target_set_task = x_target_set_task.view(-1, c, h, w) y_target_set_task = y_target_set_task.view(-1) for num_step in range(num_steps): support_loss, support_preds = self.net_forward(x=x_support_set_task, y=y_support_set_task, weights=names_weights_copy, backup_running_statistics= True if (num_step == 0) else False, training=True, num_step=num_step) names_weights_copy = self.apply_inner_loop_update(loss=support_loss, names_weights_copy=names_weights_copy, use_second_order=use_second_order, current_step_idx=num_step) if use_multi_step_loss_optimization and training_phase and epoch < self.args.multi_step_loss_num_epochs: target_loss, target_preds = self.net_forward(x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) task_losses.append(per_step_loss_importance_vectors[num_step] * target_loss) else: if num_step == (self.args.number_of_training_steps_per_iter - 1): target_loss, target_preds = self.net_forward(x=x_target_set_task, y=y_target_set_task, weights=names_weights_copy, backup_running_statistics=False, training=True, num_step=num_step) task_losses.append(target_loss) per_task_target_preds[task_id] = target_preds.detach().cpu().numpy() _, predicted = torch.max(target_preds.data, 1) #row_sums = torch.sum(target_preds, 1) #row_sums = row_sums.repeat(1, len(per_task_target_preds[task_id])) #row_sums = row_sums.reshape(target_preds.shape) #probs = torch.div(target_preds, row_sums) #print(per_task_target_preds[task_id]) #print(probs) #print(target_preds) accuracy = predicted.float().eq(y_target_set_task.data.float()).cpu().float() #print("-----") #print("accuracy: ") #print(accuracy) #print("target") #print(y_target_set_task.cpu().data.float()) #print("pred") #print(predicted.cpu().data.float()) #roc = roc_auc_score(y_target_set_task.cpu().data.float(), probs, average="macro", multi_class="ovo") roc = 0.5 #print("roc: ") #print(roc) #f1 = f1_score(y_target_set_task.cpu().data.float(), probs, average) f1 = 1 #print("f1: ") #print(f1) #print(task_losses) task_losses = torch.sum(torch.stack(task_losses)) #print(task_losses) total_losses.append(task_losses) total_accuracies.extend(accuracy) total_auc.append(roc) #print(total_auc) total_f1.append(f1) #print(total_f1) #print("-----") if not training_phase: self.classifier.restore_backup_stats() losses, metrics = self.get_across_task_loss_metrics(total_losses=total_losses, total_accuracies=total_accuracies, total_auc=total_auc, total_f1=total_f1) for idx, item in enumerate(per_step_loss_importance_vectors): losses['loss_importance_vector_{}'.format(idx)] = item.detach().cpu().numpy() return losses, metrics, per_task_target_preds def net_forward(self, x, y, weights, backup_running_statistics, training, num_step): """ A base model forward pass on some data points x. Using the parameters in the weights dictionary. Also requires boolean flags indicating whether to reset the running statistics at the end of the run (if at evaluation phase). A flag indicating whether this is the training session and an int indicating the current step's number in the inner loop. :param x: A data batch of shape b, c, h, w :param y: A data targets batch of shape b, n_classes :param weights: A dictionary containing the weights to pass to the network. :param backup_running_statistics: A flag indicating whether to reset the batch norm running statistics to their previous values after the run (only for evaluation) :param training: A flag indicating whether the current process phase is a training or evaluation. :param num_step: An integer indicating the number of the step in the inner loop. :return: the crossentropy losses with respect to the given y, the predictions of the base model. """ preds = self.classifier.forward(x=x, params=weights, training=training, backup_running_statistics=backup_running_statistics, num_step=num_step) loss = F.cross_entropy(input=preds, target=y) return loss, preds def trainable_parameters(self): """ Returns an iterator over the trainable parameters of the model. """ for param in self.parameters(): if param.requires_grad: yield param def train_forward_prop(self, data_batch, epoch): """ Runs an outer loop forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, metrics, per_task_target_preds = self.forward(data_batch=data_batch, epoch=epoch, use_second_order=self.args.second_order and epoch > self.args.first_order_to_second_order_epoch, use_multi_step_loss_optimization=self.args.use_multi_step_loss_optimization, num_steps=self.args.number_of_training_steps_per_iter, training_phase=True) return losses, metrics, per_task_target_preds def evaluation_forward_prop(self, data_batch, epoch): """ Runs an outer loop evaluation forward prop using the meta-model and base-model. :param data_batch: A data batch containing the support set and the target set input, output pairs. :param epoch: The index of the currrent epoch. :return: A dictionary of losses for the current step. """ losses, metrics, per_task_target_preds = self.forward(data_batch=data_batch, epoch=epoch, use_second_order=False, use_multi_step_loss_optimization=True, num_steps=self.args.number_of_evaluation_steps_per_iter, training_phase=False) return losses, metrics, per_task_target_preds def meta_update(self, loss): """ Applies an outer loop update on the meta-parameters of the model. :param loss: The current crossentropy loss. """ self.optimizer.zero_grad() loss.backward() if 'imagenet' in self.args.dataset_name: for name, param in self.classifier.named_parameters(): if param.requires_grad: param.grad.data.clamp_(-10, 10) # not sure if this is necessary, more experiments are needed self.optimizer.step() def run_train_iter(self, data_batch, epoch): """ Runs an outer loop update step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ epoch = int(epoch) self.scheduler.step() if self.current_epoch != epoch: self.current_epoch = epoch if not self.training: self.train() x_support_set, x_target_set, y_support_set, y_target_set = data_batch x_support_set = torch.Tensor(x_support_set).float().to(device=self.device) x_target_set = torch.Tensor(x_target_set).float().to(device=self.device) y_support_set = torch.Tensor(y_support_set).long().to(device=self.device) y_target_set = torch.Tensor(y_target_set).long().to(device=self.device) data_batch = (x_support_set, x_target_set, y_support_set, y_target_set) losses, metrics, per_task_target_preds = self.train_forward_prop(data_batch=data_batch, epoch=epoch) self.meta_update(loss=losses['loss']) losses['learning_rate'] = self.scheduler.get_lr()[0] self.optimizer.zero_grad() self.zero_grad() return losses, metrics, per_task_target_preds def run_validation_iter(self, data_batch): """ Runs an outer loop evaluation step on the meta-model's parameters. :param data_batch: input data batch containing the support set and target set input, output pairs :param epoch: the index of the current epoch :return: The losses of the ran iteration. """ if self.training: self.eval() x_support_set, x_target_set, y_support_set, y_target_set = data_batch x_support_set = torch.Tensor(x_support_set).float().to(device=self.device) x_target_set = torch.Tensor(x_target_set).float().to(device=self.device) y_support_set = torch.Tensor(y_support_set).long().to(device=self.device) y_target_set = torch.Tensor(y_target_set).long().to(device=self.device) data_batch = (x_support_set, x_target_set, y_support_set, y_target_set) losses, metrics, per_task_target_preds = self.evaluation_forward_prop(data_batch=data_batch, epoch=self.current_epoch) # losses['loss'].backward() # uncomment if you get the weird memory error # self.zero_grad() # self.optimizer.zero_grad() return losses, metrics, per_task_target_preds def save_model(self, model_save_dir, state): """ Save the network parameter state and experiment state dictionary. :param model_save_dir: The directory to store the state at. :param state: The state containing the experiment state and the network. It's in the form of a dictionary object. """ state['network'] = self.state_dict() torch.save(state, f=model_save_dir) def load_model(self, model_save_dir, model_name, model_idx): """ Load checkpoint and return the state dictionary containing the network state params and experiment state. :param model_save_dir: The directory from which to load the files. :param model_name: The model_name to be loaded from the direcotry. :param model_idx: The index of the model (i.e. epoch number or 'latest' for the latest saved model of the current experiment) :return: A dictionary containing the experiment state and the saved model parameters. """ filepath = os.path.join(model_save_dir, "{}_{}".format(model_name, model_idx)) state = torch.load(filepath) state_dict_loaded = state['network'] self.load_state_dict(state_dict=state_dict_loaded) return state
def __init__(self, device, args): """ Initializes a MAML few shot learning system :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLFewShotClassifier, self).__init__(device, args) config = AutoConfig.from_pretrained(args.pretrained_weights) config.num_labels = args.num_classes_per_set model_initialization = AutoModelForSequenceClassification.from_pretrained( args.pretrained_weights, config=config ) slow_model = MetaBERT # Init fast model state_dict = model_initialization.state_dict() config = model_initialization.config del model_initialization # Slow model self.classifier = slow_model.init_from_pretrained( state_dict, config, num_labels=args.num_classes_per_set, is_distil=self.is_distil, is_xlm=self.is_xlm, per_step_layer_norm_weights=args.per_step_layer_norm_weights, num_inner_loop_steps=args.number_of_training_steps_per_iter, device=device, ) self.classifier.to("cpu") self.classifier.train() self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=torch.device("cpu"), init_learning_rate=self.task_learning_rate, total_num_inner_loop_steps=self.args.number_of_training_steps_per_iter, use_learnable_learning_rates=self.args.learnable_per_layer_per_step_inner_loop_learning_rate, init_class_head_lr_multiplier=self.args.init_class_head_lr_multiplier, ) self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters() ) ) print("Inner Loop parameters") for key, value in self.inner_loop_optimizer.named_parameters(): print(key, value.shape) print("Outer Loop parameters") for name, param in self.named_parameters(): if param.requires_grad: print(name, param.shape, param.device, param.requires_grad) self.optimizer = Ranger( [ {"params": self.classifier.parameters(), "lr": args.meta_learning_rate}, { "params": self.inner_loop_optimizer.parameters(), "lr": args.meta_inner_optimizer_learning_rate, }, ], lr=args.meta_learning_rate, ) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs * self.args.total_iter_per_epoch, eta_min=self.args.min_learning_rate, ) self.inner_loop_optimizer.to(self.device) self.clip_value = 1.0 # gradient clipping for p in self.classifier.parameters(): if p.requires_grad: p.register_hook( lambda grad: torch.clamp(grad, -self.clip_value, self.clip_value) ) self.num_freeze_epochs = args.num_freeze_epochs if self.num_freeze_epochs > 0: self.classifier.freeze()
class MAMLFewShotClassifier(FewShotClassifier): def __init__(self, device, args): """ Initializes a MAML few shot learning system :param device: The device to use to use the model on. :param args: A namedtuple of arguments specifying various hyperparameters. """ super(MAMLFewShotClassifier, self).__init__(device, args) config = AutoConfig.from_pretrained(args.pretrained_weights) config.num_labels = args.num_classes_per_set model_initialization = AutoModelForSequenceClassification.from_pretrained( args.pretrained_weights, config=config ) slow_model = MetaBERT # Init fast model state_dict = model_initialization.state_dict() config = model_initialization.config del model_initialization # Slow model self.classifier = slow_model.init_from_pretrained( state_dict, config, num_labels=args.num_classes_per_set, is_distil=self.is_distil, is_xlm=self.is_xlm, per_step_layer_norm_weights=args.per_step_layer_norm_weights, num_inner_loop_steps=args.number_of_training_steps_per_iter, device=device, ) self.classifier.to("cpu") self.classifier.train() self.inner_loop_optimizer = LSLRGradientDescentLearningRule( device=torch.device("cpu"), init_learning_rate=self.task_learning_rate, total_num_inner_loop_steps=self.args.number_of_training_steps_per_iter, use_learnable_learning_rates=self.args.learnable_per_layer_per_step_inner_loop_learning_rate, init_class_head_lr_multiplier=self.args.init_class_head_lr_multiplier, ) self.inner_loop_optimizer.initialise( names_weights_dict=self.get_inner_loop_parameter_dict( params=self.classifier.named_parameters() ) ) print("Inner Loop parameters") for key, value in self.inner_loop_optimizer.named_parameters(): print(key, value.shape) print("Outer Loop parameters") for name, param in self.named_parameters(): if param.requires_grad: print(name, param.shape, param.device, param.requires_grad) self.optimizer = Ranger( [ {"params": self.classifier.parameters(), "lr": args.meta_learning_rate}, { "params": self.inner_loop_optimizer.parameters(), "lr": args.meta_inner_optimizer_learning_rate, }, ], lr=args.meta_learning_rate, ) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer=self.optimizer, T_max=self.args.total_epochs * self.args.total_iter_per_epoch, eta_min=self.args.min_learning_rate, ) self.inner_loop_optimizer.to(self.device) self.clip_value = 1.0 # gradient clipping for p in self.classifier.parameters(): if p.requires_grad: p.register_hook( lambda grad: torch.clamp(grad, -self.clip_value, self.clip_value) ) self.num_freeze_epochs = args.num_freeze_epochs if self.num_freeze_epochs > 0: self.classifier.freeze() def get_inner_loop_parameter_dict(self, params, adapter_only=False): """ Returns a dictionary with the parameters to use for inner loop updates. :param params: A dictionary of the network's parameters. :return: A dictionary of the parameters to use for the inner loop optimization process. """ param_dict = dict() for name, param in params: if param.requires_grad: key = ( name.replace("module.", "", 1) if name.startswith("module.") else name ) if self.args.enable_inner_loop_optimizable_ln_params: param_dict[key] = param.to(device=self.device) else: if "LayerNorm" not in key: if adapter_only: if "adapter" in key or "classifier" in key: param_dict[key] = param.to(device=self.device) else: print(key) else: param_dict[key] = param.to(device=self.device) return param_dict def apply_inner_loop_update( self, loss, names_weights_copy, use_second_order, current_step_idx, allow_unused=True, ): """ Applies an inner loop update given current step's loss, the weights to update, a flag indicating whether to use second order derivatives and the current step's index. :param loss: Current step's loss with respect to the support set. :param names_weights_copy: A dictionary with names to parameters to update. :param use_second_order: A boolean flag of whether to use second order derivatives. :param current_step_idx: Current step's index. :return: A dictionary with the updated weights (name, param) """ all_names = list(names_weights_copy.keys()) names_weights_copy = { k: v for k, v in names_weights_copy.items() if v is not None } grads = torch.autograd.grad( loss, names_weights_copy.values(), create_graph=use_second_order, allow_unused=allow_unused, ) names_grads_wrt_params = dict(zip(names_weights_copy.keys(), grads)) names_weights_copy = self.inner_loop_optimizer.update_params( names_weights_dict=names_weights_copy, names_grads_wrt_params_dict=names_grads_wrt_params, num_step=current_step_idx, ) del names_grads_wrt_params for name in all_names: if name not in names_weights_copy.keys(): names_weights_copy[name] = None return names_weights_copy def net_forward( self, x, teacher_unary, fast_model, training, num_step, return_nr_correct=False, mask=None, task_name="", ): student_logits = self.classifier( input_ids=x, attention_mask=mask, num_step=num_step, params=fast_model )[0] set_kl_loss = False if task_name in self.gold_label_tasks and self.meta_loss.lower() == "kl": set_kl_loss = True self.meta_loss = "ce" loss = self.inner_loss( student_logits, teacher_unary, return_nr_correct=return_nr_correct ) if set_kl_loss: self.meta_loss = "kl" return loss def forward( self, data_batch, epoch, use_second_order, num_steps, training_phase, ): """ Runs a forward outer loop pass on the batch of tasks using the MAML/++ framework. :param data_batch: A data batch containing the support and target sets. :param epoch: Current epoch's index :param use_second_order: A boolean saying whether to use second order derivatives. :param use_multi_step_loss_optimization: Whether to optimize on the outer loop using just the last step's target loss (True) or whether to use multi step loss which improves the stability of the system (False) :param num_steps: Number of inner loop steps. :param training_phase: Whether this is a training phase (True) or an evaluation phase (False) :return: A dictionary with the collected losses of the current outer forward propagation. """ ( x_support_set, len_support_set, x_target_set, len_target_set, y_support_set, y_target_set, teacher_names, ) = data_batch meta_batch_size = self.args.batch_size self.classifier.zero_grad() if self.num_freeze_epochs <= epoch: self.classifier.unfreeze() losses = {"loss": 0} task_accuracies = [] task_lang_logs = [] for ( task_id, ( x_support_set_task, len_support_set_task, y_support_set_task, x_target_set_task, len_target_set_task, y_target_set_task, teacher_name, ), ) in enumerate( zip( x_support_set, len_support_set, y_support_set, x_target_set, len_target_set, y_target_set, teacher_names, ) ): task_lang_log = [teacher_name, epoch] task_losses = [] # freeze and unfreeze if necessary to get correct params if epoch <= self.num_freeze_epochs: self.classifier.unfreeze() fast_weights = self.classifier.get_inner_loop_params() if epoch < self.num_freeze_epochs: self.classifier.freeze() total_task_loss = 0 x_support_set_task = x_support_set_task.squeeze() len_support_set_task = len_support_set_task.squeeze() y_support_set_task = y_support_set_task.squeeze() x_target_set_task = x_target_set_task.squeeze() len_target_set_task = len_target_set_task.squeeze() y_target_set_task = y_target_set_task.squeeze() for num_step in range(num_steps): torch.cuda.empty_cache() support_loss, is_correct = self.net_forward( x=x_support_set_task, mask=len_support_set_task, num_step=num_step, teacher_unary=y_support_set_task, fast_model=fast_weights, training=True, return_nr_correct=True, task_name=teacher_name, ) fast_weights = self.apply_inner_loop_update( loss=support_loss, names_weights_copy=fast_weights, use_second_order=use_second_order, current_step_idx=num_step, ) if num_step == (self.args.number_of_training_steps_per_iter - 1): # store support set statistics task_lang_log.append(support_loss.detach().item()) task_lang_log.append(np.mean(is_correct)) target_loss, is_correct = self.net_forward( x=x_target_set_task, mask=len_target_set_task, teacher_unary=y_target_set_task, num_step=num_step, fast_model=fast_weights, training=True, return_nr_correct=True, task_name=teacher_name, ) task_losses.append(target_loss) accuracy = np.mean(is_correct) task_accuracies.append(accuracy) # store query set statistics task_lang_log.append(target_loss.detach().item()) task_lang_log.append(accuracy) # Achieve gradient accumulation by already backpropping current loss torch.cuda.empty_cache() task_losses = torch.sum(torch.stack(task_losses)) / meta_batch_size task_losses.backward() total_task_loss += task_losses.detach().cpu().item() losses["loss"] += total_task_loss task_lang_logs.append(task_lang_log) torch.cuda.synchronize() losses["accuracy"] = np.mean(task_accuracies) if training_phase: return losses, task_lang_logs else: return losses def finetune_epoch( self, names_weights_copy, model_config, train_dataloader, dev_dataloader, best_loss, eval_every, model_save_dir, task_name, epoch, train_on_cpu=False, writer=None, ): """ Finetunes the meta-learned classifier on a dataset :param train_dataloader: Dataloader with train examples :param dev_dataloader: Dataloader with validation examples :param best_loss: best achieved loss on dev set up till now :param eval_every: eval on dev set after eval_every updates :param model_save_dir: directory to save the model to :param task_name: name of the task finetuning is performed on :param epoch: current epoch number :return: best_loss """ if train_on_cpu: self.device = torch.device("cpu") self.inner_loop_optimizer.requires_grad_(False) self.inner_loop_optimizer.eval() self.inner_loop_optimizer.to(self.device) self.classifier.to(self.device) if names_weights_copy is None: if epoch <= self.num_freeze_epochs: self.classifier.unfreeze() # # Get fast weights names_weights_copy = self.classifier.get_inner_loop_params() if epoch < self.num_freeze_epochs: self.classifier.freeze() eval_every = ( eval_every if eval_every < len(train_dataloader) else len(train_dataloader) ) if writer is not None: # create histogram of weights for param_name, param in names_weights_copy.items(): writer.add_histogram(task_name + "/" + param_name, param, 0) writer.flush() with tqdm( initial=0, total=eval_every * self.args.number_of_training_steps_per_iter ) as pbar_train: for batch_idx, batch in enumerate(train_dataloader): torch.cuda.empty_cache() batch = tuple(t.to(self.device) for t in batch) x, mask, y_true = batch for train_step in range(self.args.number_of_training_steps_per_iter): support_loss = self.net_forward( x, mask=mask, teacher_unary=y_true, num_step=train_step, fast_model=names_weights_copy, training=True, ) names_weights_copy = self.apply_inner_loop_update( loss=support_loss, names_weights_copy=names_weights_copy, use_second_order=False, current_step_idx=train_step, ) self.inner_loop_optimizer.zero_grad() pbar_train.update(1) pbar_train.set_description( "finetuning phase {} -> loss: {}".format( batch_idx * self.args.number_of_training_steps_per_iter + train_step + 1, support_loss.item(), ) ) if writer is not None: # create histogram of weights for param_name, param in names_weights_copy.items(): writer.add_histogram( task_name + "/" + param_name, param, train_step + 1 ) writer.flush() if (batch_idx + 1) % eval_every == 0: print("Evaluating model...") losses = [] is_correct_preds = [] if train_on_cpu: self.device = torch.device("cuda") self.classifier.to(self.device) with torch.no_grad(): for batch in tqdm( dev_dataloader, desc="Evaluating", leave=False, total=len(dev_dataloader), ): batch = tuple(t.to(self.device) for t in batch) x, mask, y_true = batch loss, is_correct = self.net_forward( x, mask=mask, teacher_unary=y_true, fast_model=names_weights_copy, training=False, return_nr_correct=True, num_step=train_step, ) losses.append(loss.item()) is_correct_preds.extend(is_correct.tolist()) avg_loss = np.mean(losses) accuracy = np.mean(is_correct_preds) print("Accuracy", accuracy) if avg_loss < best_loss: best_loss = avg_loss print( "New best finetuned model with loss {:.05f}".format( best_loss ) ) torch.save( names_weights_copy, os.path.join( model_save_dir, "model_finetuned_{}".format( task_name.replace("train/", "", 1) .replace("val/", "", 1) .replace("test/", "", 1) ), ), ) return names_weights_copy, best_loss, avg_loss, accuracy