def run_meta_iteration(i_iter): # In each meta-iteration we draw a meta-batch of several tasks # Then we take a grad step with theta. # Generate the data sets of the training-tasks for meta-batch: mb_data_loaders = task_generator.create_meta_batch( prm, meta_batch_size, meta_split='meta_train') # For each task, prepare an iterator to generate training batches: mb_iterators = [ iter(mb_data_loaders[ii]['train']) for ii in range(meta_batch_size) ] # Get objective based on tasks in meta-batch: total_objective, info = meta_step(prm, model, mb_data_loaders, mb_iterators, loss_criterion) # Take gradient step with the meta-parameters (theta) based on validation data: grad_step(total_objective, meta_optimizer, lr_schedule, prm.lr, i_iter) # Print status: log_interval = 5 if (i_iter) % log_interval == 0: batch_acc = info['correct_count'] / info['sample_count'] print( cmn.status_string(i_iter, n_iterations, 1, 1, batch_acc, total_objective.data[0]))
def run_train_epoch(i_epoch): # log_interval = 500 post_model.train() train_iterators = iter(train_loader['train']) for batch_idx, batch_data in enumerate(train_loader['train']): task_loss, info = get_objective(prior_model, prm, [train_loader], object.feval, [train_iterators], [post_model], loss_criterion, 1, [0]) grad_step(task_loss[0], post_model, loss_criterion, optimizer, prm, train_iterators, train_loader['train'], lr_schedule, prm.optim_args['lr'], i_epoch) # for log_var in post_model.parameters(): # if log_var.requires_grad is False: # log_var.data = log_var.data - (i_epoch + 1) * math.log(1 + prm.gamma1) # Print status: log_interval = 10 if (batch_idx) % log_interval == 0: batch_acc = info['correct_count'] / info['sample_count'] print( cmn.status_string(i_epoch, prm.n_meta_train_epochs, batch_idx, n_batches, batch_acc) + ' Empiric-Loss: {:.4f}'.format(info['avg_empirical_loss']))
def run_train_epoch(i_epoch): log_interval = 500 post_model.train() for batch_idx, batch_data in enumerate(train_loader): correct_count = 0 sample_count = 0 # Monte-Carlo iterations: n_MC = prm.n_MC task_empirical_loss = 0 task_complexity = 0 for i_MC in range(n_MC): # get batch: inputs, targets = data_gen.get_batch_vars(batch_data, prm) # Calculate empirical loss: outputs = post_model(inputs) curr_empirical_loss = loss_criterion(outputs, targets) curr_empirical_loss, curr_complexity = get_bayes_task_objective( prm, prior_model, post_model, n_train_samples, curr_empirical_loss, noised_prior=False) task_empirical_loss += (1 / n_MC) * curr_empirical_loss task_complexity += (1 / n_MC) * curr_complexity correct_count += count_correct(outputs, targets) sample_count += inputs.size(0) # Total objective: total_objective = task_empirical_loss + task_complexity # Take gradient step with the posterior: grad_step(total_objective, optimizer, lr_schedule, prm.lr, i_epoch) # Print status: if batch_idx % log_interval == 0: batch_acc = correct_count / sample_count print( cmn.status_string(i_epoch, prm.n_meta_test_epochs, batch_idx, n_batches, batch_acc, total_objective.item()) + ' Empiric Loss: {:.4}\t Intra-Comp. {:.4}'.format( task_empirical_loss.item(), task_complexity.item())) data_objective.append(total_objective.item()) data_accuracy.append(batch_acc) data_emp_loss.append(task_empirical_loss.item()) data_task_comp.append(task_complexity.item()) return total_objective.item()
def run_train_epoch(i_epoch, i_step = 0): # For each task, prepare an iterator to generate training batches: train_iterators = [iter(data_loaders[ii]['train']) for ii in range(n_train_tasks)] # The task order to take batches from: # The meta-batch will be balanced - i.e, each task will appear roughly the same number of times # note: if some tasks have less data that other tasks - it may be sampled more than once in an epoch task_order = [] task_ids_list = list(range(n_train_tasks)) for i_batch in range(n_batches_per_task): random.shuffle(task_ids_list) task_order += task_ids_list # Note: this method ensures each training sample in each task is drawn in each epoch. # If all the tasks have the same number of sample, then each sample is drawn exactly once in an epoch. # random.shuffle(task_ids_list) # --############ TEMP # ----------- meta-batches loop (batches of tasks) -----------------------------------# # each meta-batch includes several tasks # we take a grad step with theta after each meta-batch meta_batch_starts = list(range(0, len(task_order), prm.meta_batch_size)) n_meta_batches = len(meta_batch_starts) for i_meta_batch in range(n_meta_batches): meta_batch_start = meta_batch_starts[i_meta_batch] task_ids_in_meta_batch = task_order[meta_batch_start: (meta_batch_start + prm.meta_batch_size)] # meta-batch size may be less than prm.meta_batch_size at the last one # note: it is OK if some tasks appear several times in the meta-batch mb_data_loaders = [data_loaders[task_id] for task_id in task_ids_in_meta_batch] mb_iterators = [train_iterators[task_id] for task_id in task_ids_in_meta_batch] mb_posteriors_models = [posteriors_models[task_id] for task_id in task_ids_in_meta_batch] # prior_weight_steps = 10000 # # prior_weight = 1 - math.exp(-i_step/prior_weight_steps) # prior_weight = min(i_step / prior_weight_steps, 1.0) i_step += 1 # Get objective based on tasks in meta-batch: total_objective, info = get_objective(prior_model, prm, mb_data_loaders, mb_iterators, mb_posteriors_models, loss_criterion, n_train_tasks) # Take gradient step with the shared prior and all tasks' posteriors: grad_step(total_objective, all_optimizer, lr_schedule, prm.lr, i_epoch) # Print status: log_interval = 200 if i_meta_batch % log_interval == 0: batch_acc = info['correct_count'] / info['sample_count'] print(cmn.status_string(i_epoch, prm.n_meta_train_epochs, i_meta_batch, n_meta_batches, batch_acc, total_objective.item()) + ' Empiric-Loss: {:.4}\t Task-Comp. {:.4}\t Meta-Comp.: {:.4}\t'. format(info['avg_empirical_loss'], info['avg_intra_task_comp'], info['meta_comp'])) # end meta-batches loop return i_step
def run_train_epoch(i_epoch): # # Adjust randomness (eps_std) # if hasattr(prm, 'use_randomness_schedeule') and prm.use_randomness_schedeule: # if i_epoch > prm.randomness_full_epoch: # eps_std = 1.0 # elif i_epoch > prm.randomness_init_epoch: # eps_std = (i_epoch - prm.randomness_init_epoch) / (prm.randomness_full_epoch - prm.randomness_init_epoch) # else: # eps_std = 0.0 # turn off randomness # post_model.set_eps_std(eps_std) # post_model.set_eps_std(0.00) # debug complexity_term = 0 post_model.train() for batch_idx, batch_data in enumerate(train_loader): # Monte-Carlo iterations: empirical_loss = 0 n_MC = prm.n_MC for i_MC in range(n_MC): # get batch: inputs, targets = data_gen.get_batch_vars(batch_data, prm) # calculate objective: outputs = post_model(inputs) empirical_loss_c = loss_criterion(outputs, targets) empirical_loss += (1 / n_MC) * empirical_loss_c # complexity/prior term: if prior_model: empirical_loss, complexity_term = get_bayes_task_objective( prm, prior_model, post_model, n_train_samples, empirical_loss) else: complexity_term = 0.0 # Total objective: objective = empirical_loss + complexity_term # Take gradient step: grad_step(objective, optimizer, lr_schedule, prm.lr, i_epoch) # Print status: log_interval = 500 if batch_idx % log_interval == 0: batch_acc = correct_rate(outputs, targets) print( cmn.status_string(i_epoch, prm.num_epochs, batch_idx, n_batches, batch_acc, objective.data[0]) + ' Loss: {:.4}\t Comp.: {:.4}'.format( get_value(empirical_loss), get_value(complexity_term)))
def run_train_epoch(i_epoch): # For each task, prepare an iterator to generate training batches: train_iterators = [ iter(train_data_loaders[ii]['train']) for ii in range(n_tasks) ] # The task order to take batches from: task_order = [] task_ids_list = list(range(n_tasks)) for i_batch in range(n_batches_per_task): random.shuffle(task_ids_list) task_order += task_ids_list # each meta-batch includes several tasks # we take a grad step with theta after each meta-batch meta_batch_starts = list(range(0, len(task_order), prm.meta_batch_size)) n_meta_batches = len(meta_batch_starts) # ----------- meta-batches loop (batches of tasks) -----------------------------------# for i_meta_batch in range(n_meta_batches): meta_batch_start = meta_batch_starts[i_meta_batch] task_ids_in_meta_batch = task_order[meta_batch_start:( meta_batch_start + prm.meta_batch_size)] n_tasks_in_batch = len( task_ids_in_meta_batch ) # it may be less than prm.meta_batch_size at the last one # note: it is OK if some task appear several times in the meta-batch mb_data_loaders = [ train_data_loaders[task_id] for task_id in task_ids_in_meta_batch ] mb_iterators = [ train_iterators[task_id] for task_id in task_ids_in_meta_batch ] # Get objective based on tasks in meta-batch: total_objective, info = meta_step(prm, model, mb_data_loaders, mb_iterators, loss_criterion) # Take gradient step with the meta-parameters (theta) based on validation data: grad_step(total_objective, meta_optimizer, lr_schedule, prm.lr, i_epoch) # Print status: log_interval = 200 if i_meta_batch % log_interval == 0: batch_acc = info['correct_count'] / info['sample_count'] print( cmn.status_string(i_epoch, num_epochs, i_meta_batch, n_meta_batches, batch_acc, total_objective.item()))
def run_train_epoch(i_epoch, log_mat): post_model.train() for batch_idx, batch_data in enumerate(train_loader): # get batch data: inputs, targets = data_gen.get_batch_vars(batch_data, prm) batch_size = inputs.shape[0] # Monte-Carlo iterations: avg_empiric_loss = torch.zeros(1, device=prm.device) n_MC = prm.n_MC for i_MC in range(n_MC): # calculate objective: outputs = post_model(inputs) avg_empiric_loss_curr = (1 / batch_size) * loss_criterion( outputs, targets) avg_empiric_loss += (1 / n_MC) * avg_empiric_loss_curr # complexity/prior term: if prior_model: complexity_term = get_task_complexity(prm, prior_model, post_model, n_train_samples, avg_empiric_loss) else: complexity_term = torch.zeros(1, device=prm.device) # Total objective: objective = avg_empiric_loss + complexity_term # Take gradient step: grad_step(objective, optimizer, lr_schedule, prm.lr, i_epoch) # Print status: log_interval = 1000 if batch_idx % log_interval == 0: batch_acc = correct_rate(outputs, targets) print( cmn.status_string(i_epoch, prm.num_epochs, batch_idx, n_batches, batch_acc, objective.item()) + ' Loss: {:.4}\t Comp.: {:.4}'.format( avg_empiric_loss.item(), complexity_term.item())) # End batch loop # save results for epochs-figure: if figure_flag and (i_epoch % prm.log_figure['interval_epochs'] == 0): save_result_for_figure(post_model, prior_model, data_loader, prm, log_mat, i_epoch)
def run_train_epoch(i_epoch): log_interval = 500 model.train() for batch_idx, batch_data in enumerate(train_loader): # get batch: inputs, targets = data_gen.get_batch_vars(batch_data, prm) # Calculate loss: outputs = model(inputs) loss = loss_criterion(outputs, targets) # Take gradient step: grad_step(loss, optimizer, lr_schedule, prm.lr, i_epoch) # Print status: if batch_idx % log_interval == 0: batch_acc = correct_rate(outputs, targets) print( cmn.status_string(i_epoch, prm.num_epochs, batch_idx, n_batches, batch_acc, get_value(loss)))
def run_meta_iteration(i_iter, prior_model, task_generator, prm): # In each meta-iteration we draw a meta-batch of several tasks # Then we take a grad step with prior. # Unpack parameters: optim_func, optim_args, lr_schedule = \ prm.optim_func, prm.optim_args, prm.lr_schedule # Loss criterion loss_criterion = get_loss_criterion(prm.loss_type) meta_batch_size = prm.meta_batch_size n_inner_steps = prm.n_inner_steps n_meta_iterations = prm.n_meta_train_epochs # Generate the data sets of the training-tasks for meta-batch: mb_data_loaders = task_generator.create_meta_batch(prm, meta_batch_size, meta_split='meta_train') # For each task, prepare an iterator to generate training batches: mb_iterators = [ iter(mb_data_loaders[ii]['train']) for ii in range(meta_batch_size) ] # The posteriors models will adjust to new tasks in eacxh meta-batch # Create posterior models for each task: posteriors_models = [get_model(prm) for _ in range(meta_batch_size)] init_from_prior = True if init_from_prior: for post_model in posteriors_models: post_model.load_state_dict(prior_model.state_dict()) # Gather all tasks posterior params: all_post_param = sum([ list(posterior_model.parameters()) for posterior_model in posteriors_models ], []) # Create optimizer for all parameters (posteriors + prior) prior_params = list(prior_model.parameters()) all_params = all_post_param + prior_params all_optimizer = optim_func(all_params, **optim_args) # all_optimizer = optim_func(prior_params, **optim_args) ## DeBUG test_acc_avg = 0.0 for i_inner_step in range(n_inner_steps): # Get objective based on tasks in meta-batch: total_objective, info = get_objective(prior_model, prm, mb_data_loaders, mb_iterators, posteriors_models, loss_criterion, prm.n_train_tasks) # Take gradient step with the meta-parameters (theta) based on validation data: grad_step(total_objective, all_optimizer, lr_schedule, prm.lr, i_iter) # Print status: log_interval = 20 if (i_inner_step) % log_interval == 0: batch_acc = info['correct_count'] / info['sample_count'] print( cmn.status_string(i_iter, n_meta_iterations, i_inner_step, n_inner_steps, batch_acc, total_objective.data[0]) + ' Empiric-Loss: {:.4}\t Task-Comp. {:.4}\t'.format( info['avg_empirical_loss'], info['avg_intra_task_comp'])) # Print status = on test set of meta-batch: log_interval_eval = 10 if (i_iter) % log_interval_eval == 0 and i_iter > 0: test_acc_avg = run_test(mb_data_loaders, posteriors_models, loss_criterion, prm) print('Meta-iter: {} \t Meta-Batch Test Acc: {:1.3}\t'.format( i_iter, test_acc_avg)) # End of inner steps return prior_model, posteriors_models, test_acc_avg
def run_train_epoch(i_epoch): log_interval = 500 post_model.train() train_info = {} train_info["task_comp"] = 0.0 train_info["total_loss"] = 0.0 cnt = 0 for batch_idx, batch_data in enumerate(train_loader): cnt += 1 correct_count = 0 sample_count = 0 # Monte-Carlo iterations: n_MC = prm.n_MC task_empirical_loss = 0 task_complexity = 0 for i_MC in range(n_MC): # get batch: inputs, targets = data_gen.get_batch_vars(batch_data, prm) # Calculate empirical loss: outputs = post_model(inputs) curr_empirical_loss = loss_criterion(outputs, targets) #hyper_kl = 0 when testing curr_empirical_loss, curr_complexity, task_info = get_bayes_task_objective( prm, prior_model, post_model, n_train_samples, curr_empirical_loss, noised_prior=False) task_empirical_loss += (1 / n_MC) * curr_empirical_loss task_complexity += (1 / n_MC) * curr_complexity correct_count += count_correct(outputs, targets) sample_count += inputs.size(0) # Total objective: total_objective = task_empirical_loss + prm.task_complex_w * task_complexity train_info["task_comp"] += task_complexity.data[0] train_info["total_loss"] += total_objective.data[0] # Take gradient step with the posterior: grad_step(total_objective, optimizer, lr_schedule, prm.lr, i_epoch) # Print status: if batch_idx % log_interval == 0: batch_acc = correct_count / sample_count write_to_log( cmn.status_string(i_epoch, prm.n_meta_test_epochs, batch_idx, n_batches, batch_acc, total_objective.data[0]) + ' Empiric Loss: {:.4}\t Intra-Comp. {:.4}, w_kld {:.4}, b_kld {:.4}' .format(task_empirical_loss.data[0], task_complexity.data[0], task_info["w_kld"], task_info["b_kld"]), prm) train_info["task_comp"] /= cnt train_info["total_loss"] /= cnt return train_info
def run_meta_iteration(i_iter, prior_model, task_generator, prm): # In each meta-iteration we draw a meta-batch of several tasks # Then we take a grad step with prior. # Unpack parameters: optim_func, optim_args, lr_schedule = \ prm.optim_func, prm.optim_args, prm.lr_schedule # Loss criterion loss_criterion = get_loss_criterion(prm.loss_type) meta_batch_size = prm.meta_batch_size n_inner_steps = prm.n_inner_steps n_meta_iterations = prm.n_meta_train_epochs # Generate the data sets of the training-tasks for meta-batch: mb_data_loaders = task_generator.create_meta_batch(prm, meta_batch_size, meta_split='meta_train') # For each task, prepare an iterator to generate training batches: mb_iterators = [ iter(mb_data_loaders[ii]['train']) for ii in range(meta_batch_size) ] # The posteriors models will adjust to new tasks in eacxh meta-batch # Create posterior models for each task: posteriors_models = [get_model(prm) for _ in range(meta_batch_size)] init_from_prior = True if init_from_prior: for post_model in posteriors_models: post_model.load_state_dict(prior_model.state_dict()) # # Gather all tasks posterior params: # all_post_param = sum([list(posterior_model.parameters()) for posterior_model in posteriors_models], []) # # # Create optimizer for all parameters (posteriors + prior) # prior_params = list(prior_model.parameters()) # all_params = all_post_param + prior_params # all_optimizer = optim_func(all_params, **optim_args) # # all_optimizer = optim_func(prior_params, **optim_args) ## DeBUG prior_params = filter(lambda p: p.requires_grad, prior_model.parameters()) prior_optimizer = optim_func(prior_params, optim_args) if prior_optimizer.param_groups[0]['params'][0].grad is None: object.grad_init(prm, prior_model, loss_criterion, iter(mb_data_loaders[0]['train']), mb_data_loaders[0]['train'], prior_optimizer) # all_params = all_post_param + prior_params all_posterior_optimizers = [ optim_func( filter(lambda p: p.requires_grad, posteriors_models[i].parameters()), optim_args) for i in range(meta_batch_size) ] test_acc_avg = 0.0 for i_inner_step in range(n_inner_steps): # Get objective based on tasks in meta-batch: # total_objective, info = get_objective(prior_model, prm, mb_data_loaders, mb_iterators, # posteriors_models, loss_criterion, prm.n_train_tasks) task_loss_list, info = get_objective(prior_model, prm, mb_data_loaders, object.feval, mb_iterators, posteriors_models, loss_criterion, meta_batch_size, range(meta_batch_size)) # Take gradient step with the meta-parameters (theta) based on validation data: # grad_step(total_objective, all_optimizer, lr_schedule, prm.lr, i_iter) prior_optimizer.zero_grad() for i_task in range(meta_batch_size): if isinstance(task_loss_list[i_task], int): continue grad_step(task_loss_list[i_task], posteriors_models[i_task], loss_criterion, all_posterior_optimizers[i_task], prm, mb_iterators[i_task], mb_data_loaders[i_task]['train'], lr_schedule, prm.lr, i_iter) # if i_meta_batch==n_meta_batches-1: prior_get_grad(prior_optimizer, all_posterior_optimizers[i_task]) # prior_grad_step(prior_model, prior_optimizer, all_posterior_optimizers, prm.meta_batch_size, prm, # lr_schedule, prm.prior_lr, i_epoch) prior_grad_step(prior_optimizer, prm.meta_batch_size, prm, prm.prior_lr_schedule, prm.prior_lr, i_iter) # Print status: log_interval = 1 # if (i_inner_step) % log_interval == 0: # batch_acc = info['correct_count'] / info['sample_count'] # print(cmn.status_string(i_iter, n_meta_iterations, i_inner_step, n_inner_steps, batch_acc, total_objective.data[0]) + # ' Empiric-Loss: {:.4}\t Task-Comp. {:.4}\t'. # format(info['avg_empirical_loss'], info['avg_intra_task_comp'])) if (i_inner_step) % log_interval == 0: batch_acc = info['correct_count'] / info['sample_count'] print( cmn.status_string(i_iter, prm.n_meta_train_epochs, i_inner_step, n_inner_steps, batch_acc) + ' Empiric-Loss: {:.4f}'.format(info['avg_empirical_loss'])) # Print status = on test set of meta-batch: log_interval_eval = 1 if (i_iter) % log_interval_eval == 0: test_acc_avg = run_test(mb_data_loaders, posteriors_models, loss_criterion, prm) print('Meta-iter: {} \t Meta-Batch Test Acc: {:1.3}\t'.format( i_iter, test_acc_avg)) # End of inner steps return prior_model, posteriors_models, test_acc_avg
def run_train_epoch(i_epoch): # optim_args['L'] = L-5*i_epoch # For each task, prepare an iterator to generate training batches: train_iterators = [ iter(data_loaders[ii]['train']) for ii in range(n_train_tasks) ] # The task order to take batches from: # The meta-batch will be balanced - i.e, each task will appear roughly the same number of times # note: if some tasks have less data that other tasks - it may be sampled more than once in an epoch task_order = [] task_ids_list = list(range(n_train_tasks)) for i_batch in range(n_batches_per_task): random.shuffle(task_ids_list) task_order += task_ids_list # Note: this method ensures each training sample in each task is drawn in each epoch. # If all the tasks have the same number of sample, then each sample is drawn exactly once in an epoch. # ----------- meta-batches loop (batches of tasks) -----------------------------------# # each meta-batch includes several tasks # we take a grad step with theta after each meta-batch # meta_batch_starts = list(range(0, len(task_order), n_train_tasks)) meta_batch_starts = list(range(0, len(task_order), prm.meta_batch_size)) n_meta_batches = len(meta_batch_starts) for i_meta_batch in range(n_meta_batches): meta_batch_start = meta_batch_starts[i_meta_batch] task_ids_in_meta_batch = task_order[meta_batch_start:( meta_batch_start + prm.meta_batch_size)] # meta-batch size may be less than prm.meta_batch_size at the last one # note: it is OK if some tasks appear several times in the meta-batch mb_data_loaders = [ data_loaders[task_id] for task_id in task_ids_in_meta_batch ] mb_iterators = [ train_iterators[task_id] for task_id in task_ids_in_meta_batch ] mb_posteriors_models = [ posteriors_models[task_id] for task_id in task_ids_in_meta_batch ] #task_loss_list, info = get_objective(prior_model, prm, mb_data_loaders, object.feval, # mb_iterators, mb_posteriors_models, loss_criterion, n_train_tasks, task_ids_in_meta_batch) # Take gradient step with the shared prior and all tasks' posteriors: # for i_task in range(n_train_tasks): # prior_optimizer.zero_grad() #for i_task in range(n_train_tasks): # if isinstance(task_loss_list[i_task], int): # continue # grad_step(task_loss_list[i_task], posteriors_models[i_task], loss_criterion, all_posterior_optimizers[i_task], prm, # train_iterators[i_task], data_loaders[i_task]['train'], lr_schedule, prm.lr, i_epoch) #task_loss_list, info = get_objective(prior_model, prm, mb_data_loaders, object.feval, # mb_iterators, mb_posteriors_models, loss_criterion, n_train_tasks, task_ids_in_meta_batch) # Take gradient step with the shared prior and all tasks' posteriors: # for i_task in range(n_train_tasks): # prior_optimizer.zero_grad() #for i_task in range(n_train_tasks): # if isinstance(task_loss_list[i_task], int): # continue # grad_step(task_loss_list[i_task], posteriors_models[i_task], loss_criterion, all_posterior_optimizers[i_task], prm, # train_iterators[i_task], data_loaders[i_task]['train'], lr_schedule, prm.lr, i_epoch) # Get objective based on tasks in meta-batch: task_loss_list, info = get_objective(prior_model, prm, mb_data_loaders, object.feval, mb_iterators, mb_posteriors_models, loss_criterion, n_train_tasks, task_ids_in_meta_batch) # Take gradient step with the shared prior and all tasks' posteriors: # for i_task in range(n_train_tasks): prior_optimizer.zero_grad() for i_task in range(n_train_tasks): if isinstance(task_loss_list[i_task], int): continue grad_step(task_loss_list[i_task], posteriors_models[i_task], loss_criterion, all_posterior_optimizers[i_task], prm, train_iterators[i_task], data_loaders[i_task]['train'], lr_schedule, prm.lr, i_epoch) # if i_meta_batch==n_meta_batches-1: # if i_epoch == prm.n_meta_train_epochs-1: prior_get_grad(prior_optimizer, all_posterior_optimizers[i_task]) # task_loss_list, info = get_objective(prior_model, prm, mb_data_loaders, object.feval, # mb_iterators, mb_posteriors_models, loss_criterion, n_train_tasks, task_ids_in_meta_batch) # prior_grad_step(prior_optimizer, prm.meta_batch_size, prm,prm.prior_lr_schedule, prm.prior_lr, i_epoch) prior_updates(prior_optimizer, n_train_tasks, prm) for post_model in posteriors_models: post_model.load_state_dict(prior_model.state_dict()) # Print status: log_interval = 10 if (i_meta_batch) % log_interval == 0: batch_acc = info['correct_count'] / info['sample_count'] print( cmn.status_string(i_epoch, prm.n_meta_train_epochs, i_meta_batch, n_meta_batches, batch_acc) + ' Empiric-Loss: {:.4f}'.format(info['avg_empirical_loss']))
def run_train_epoch(i_epoch): log_interval = 500 post_model.train() for batch_idx, batch_data in enumerate(train_loader): # get batch data: inputs, targets = data_gen.get_batch_vars(batch_data, prm) batch_size = inputs.shape[0] correct_count = 0 sample_count = 0 # Monte-Carlo iterations: n_MC = prm.n_MC avg_empiric_loss = 0 complexity_term = 0 for i_MC in range(n_MC): # Calculate empirical loss: outputs = post_model(inputs) avg_empiric_loss_curr = (1 / batch_size) * loss_criterion( outputs, targets) # complexity_curr = get_task_complexity(prm, prior_model, post_model, # n_train_samples, avg_empiric_loss_curr) avg_empiric_loss += (1 / n_MC) * avg_empiric_loss_curr # complexity_term += (1 / n_MC) * complexity_curr correct_count += count_correct(outputs, targets) sample_count += inputs.size(0) # end monte-carlo loop complexity_term = get_task_complexity(prm, prior_model, post_model, n_train_samples, avg_empiric_loss) # Approximated total objective (for current batch): if prm.complexity_type == 'Variational_Bayes': # note that avg_empiric_loss_per_task is estimated by an average over batch samples, # but its weight in the objective should be considered by how many samples there are total in the task total_objective = avg_empiric_loss * ( n_train_samples) + complexity_term else: total_objective = avg_empiric_loss + complexity_term # Take gradient step with the posterior: grad_step(total_objective, optimizer, lr_schedule, prm.lr, i_epoch) # Print status: if batch_idx % log_interval == 0: batch_acc = correct_count / sample_count print( cmn.status_string(i_epoch, prm.n_meta_test_epochs, batch_idx, n_batches, batch_acc, total_objective.item()) + ' Empiric Loss: {:.4}\t Intra-Comp. {:.4}'.format( avg_empiric_loss.item(), complexity_term.item()))