wasserstein_ensemble.py

# import ot
# import torch
# import numpy as np
# # import routines
# # from model import get_model_from_name
# # import utils
# from ground_metric import GroundMetric
# import math
# import sys
# # import compute_activations
# import modelset as mynetwork
# torch.manual_seed(0)
# torch.cuda.manual_seed(0)
# np.random.seed(0)
#
#
# def cost_matrix(x, y, p=2):
#     "Returns the matrix of $|x_i-y_j|^p$."
#     x_col = x.unsqueeze(1)
#     y_lin = y.unsqueeze(0)
#     c = torch.sum((torch.abs(x_col - y_lin)) ** p, 2)
#     return c
#
# def get_histogram(args, idx, cardinality, layer_name=None,  return_numpy = True, float64=False):
#
#     if not args.unbalanced:
# #             print("returns a uniform measure of cardinality: ", cardinality)
#         return np.ones(cardinality)/cardinality
#     else:
#         return np.ones(cardinality)
#
# def get_wassersteinized_layers_modularized(args, networks, eps=1e-7, test_loader=None,pp1=None):
#     '''
#     Two neural networks that have to be averaged in geometric manner (i.e. layerwise).
#     The 1st network is aligned with respect to the other via wasserstein distance.
#     Also this assumes that all the layers are either fully connected or convolutional *(with no bias)*
#
#     :param networks: list of networks
#     :param activations: If not None, use it to build the activation histograms.
#     Otherwise assumes uniform distribution over neurons in a layer.
#     :return: list of layer weights 'wassersteinized'
#     '''
#
#     # simple_model_0, simple_model_1 = networks[0], networks[1]
#     # simple_model_0 = get_trained_model(0, model='simplenet')
#     # simple_model_1 = get_trained_model(1, model='simplenet')
#
#     avg_aligned_layers = []
#     # cumulative_T_var = None
#     T_var = None
#     # print(list(networks[0].parameters()))
#     previous_layer_shape = None
#     ground_metric_object = GroundMetric(args)
#
#     if args.eval_aligned:
#         model0_aligned_layers = []
#
#     if args.gpu_id==-1:
#         device = torch.device('cpu')
#     else:
#         device = torch.device('cuda:{}'.format(args.gpu_id))
#
#
#     num_layers = len(list(zip(networks[0].parameters(), networks[1].parameters())))
#     for idx, ((layer0_name, fc_layer0_weight), (layer1_name, fc_layer1_weight)) in \
#             enumerate(zip(networks[0].named_parameters(), networks[1].named_parameters())):
#         if 'running' in layer0_name or 'batch' in layer0_name:continue
#         assert fc_layer0_weight.shape == fc_layer1_weight.shape
# #         print("Previous layer shape is ", previous_layer_shape)
#         previous_layer_shape = fc_layer1_weight.shape
#
#         mu_cardinality = fc_layer0_weight.shape[0]
#         nu_cardinality = fc_layer1_weight.shape[0]
#
#         # mu = np.ones(fc_layer0_weight.shape[0])/fc_layer0_weight.shape[0]
#         # nu = np.ones(fc_layer1_weight.shape[0])/fc_layer1_weight.shape[0]
#
#         layer_shape = fc_layer0_weight.shape
#         if len(layer_shape) > 2:
#             is_conv = True
#             # For convolutional layers, it is (#out_channels, #in_channels, height, width)
#             fc_layer0_weight_data = fc_layer0_weight.data.view(fc_layer0_weight.shape[0], fc_layer0_weight.shape[1], -1)
# #             if pp1:
# #                 pp1[idx] = pp1[idx].view(fc_layer0_weight.shape[0], fc_layer0_weight.shape[1], -1)
#             fc_layer1_weight_data = fc_layer1_weight.data.view(fc_layer1_weight.shape[0], fc_layer1_weight.shape[1], -1)
#         else:
#             is_conv = False
#             fc_layer0_weight_data = fc_layer0_weight.data
#             fc_layer1_weight_data = fc_layer1_weight.data
#
#         if idx == 0:
#             if is_conv:
#                 M = ground_metric_object.process(fc_layer0_weight_data.view(fc_layer0_weight_data.shape[0], -1),
#                                 fc_layer1_weight_data.view(fc_layer1_weight_data.shape[0], -1))
#                 # M = cost_matrix(fc_layer0_weight_data.view(fc_layer0_weight_data.shape[0], -1),
#                 #                 fc_layer1_weight_data.view(fc_layer1_weight_data.shape[0], -1))
#             else:
#                 # print("layer data is ", fc_layer0_weight_data, fc_layer1_weight_data)
#                 M = ground_metric_object.process(fc_layer0_weight_data, fc_layer1_weight_data)
#                 # M = cost_matrix(fc_layer0_weight, fc_layer1_weight)
#
#             aligned_wt = fc_layer0_weight_data
#         else:
#
# #             print("shape of layer: model 0", fc_layer0_weight_data.shape)
# #             print("shape of layer: model 1", fc_layer1_weight_data.shape)
# #             print("shape of previous transport map", T_var.shape)
#
#             # aligned_wt = None, this caches the tensor and causes OOM
#             if is_conv:
#
#                 T_var_conv = T_var.unsqueeze(0).repeat(fc_layer0_weight_data.shape[2], 1, 1)# (#out_channels, #in_channels, height, width)
#                 aligned_wt = torch.bmm(fc_layer0_weight_data.permute(2, 0, 1), T_var_conv).permute(1, 2, 0)
# #                 if pp1!=None:
#
#
#                 M = ground_metric_object.process(
#                     aligned_wt.contiguous().view(aligned_wt.shape[0], -1),
#                     fc_layer1_weight_data.view(fc_layer1_weight_data.shape[0], -1)
#                 )
#             else:
#                 if fc_layer0_weight.data.shape[1] != T_var.shape[0]:
#                     # Handles the switch from convolutional layers to fc layers
#                     fc_layer0_unflattened = fc_layer0_weight.data.view(fc_layer0_weight.shape[0], T_var.shape[0], -1).permute(2, 0, 1)
#                     aligned_wt = torch.bmm(
#                         fc_layer0_unflattened,
#                         T_var.unsqueeze(0).repeat(fc_layer0_unflattened.shape[0], 1, 1)
#                     ).permute(1, 2, 0)
#                     aligned_wt = aligned_wt.contiguous().view(aligned_wt.shape[0], -1)
#                 else:
#                     # print("layer data (aligned) is ", aligned_wt, fc_layer1_weight_data)
#                     aligned_wt = torch.matmul(fc_layer0_weight.data, T_var)
#                 # M = cost_matrix(aligned_wt, fc_layer1_weight)
#                 M = ground_metric_object.process(aligned_wt, fc_layer1_weight)
# #                 print("ground metric is ", M)
#             if args.skip_last_layer and idx == (num_layers - 1):
# #                 print("Simple averaging of last layer weights. NO transport map needs to be computed")
#                 if args.ensemble_step != 0.5:
#                     avg_aligned_layers.append(aligned_wt)
# #                     avg_aligned_layers.append((1 - args.ensemble_step) * aligned_wt +
# #                                           args.ensemble_step * fc_layer1_weight)
#                 else:
#                     avg_aligned_layers.append(aligned_wt)
# #                     avg_aligned_layers.append((aligned_wt + fc_layer1_weight)/2)
# #                 return avg_aligned_layers
#                 if pp1!=None:
#                     return avg_aligned_layers,pp1
#                 else:
#                     return avg_aligned_layers
#
#         if args.importance is None or (idx == num_layers -1):
#             mu = get_histogram(args, 0, mu_cardinality, layer0_name)
#             nu = get_histogram(args, 1, nu_cardinality, layer1_name)
#         else:
#             # mu = _get_neuron_importance_histogram(args, aligned_wt, is_conv)
#             mu = _get_neuron_importance_histogram(args, fc_layer0_weight_data, is_conv)
#             nu = _get_neuron_importance_histogram(args, fc_layer1_weight_data, is_conv)
#             print(mu, nu)
#             assert args.proper_marginals
#
#         cpuM = M.data.cpu().numpy()
#         if args.exact:
#             T = ot.emd(mu, nu, cpuM)
# #             print("*******************************",mu.shape,nu.shape,cpuM.shape)
# #             print(nu)
#         else:
#             T = ot.bregman.sinkhorn(mu, nu, cpuM, reg=args.reg)
#         # T = ot.emd(mu, nu, log_cpuM)
# #         print("************",str(idx),"*****",T.shape)
#
#         if args.gpu_id!=-1:
#             T_var = torch.from_numpy(T).cuda(args.gpu_id).float()
# #             if pp1:
# #                 pp1[idx+1] = torch.bmm(pp1[idx+1], T_var_conv)
#         else:
#             T_var = torch.from_numpy(T).float()
#
#         # torch.set_printoptions(profile="full")
# #         print("the transport map is ", T_var)
#         # torch.set_printoptions(profile="default")
#
#         if args.correction:
#             if not args.proper_marginals:
#                 # think of it as m x 1, scaling weights for m linear combinations of points in X
#                 if args.gpu_id != -1:
#                     # marginals = torch.mv(T_var.t(), torch.ones(T_var.shape[0]).cuda(args.gpu_id))  # T.t().shape[1] = T.shape[0]
#                     marginals = torch.ones(T_var.shape[0]).cuda(args.gpu_id) / T_var.shape[0]
#                 else:
#                     # marginals = torch.mv(T_var.t(),
#                     #                      torch.ones(T_var.shape[0]))  # T.t().shape[1] = T.shape[0]
#                     marginals = torch.ones(T_var.shape[0]) / T_var.shape[0]
#                 marginals = torch.diag(1.0/(marginals + eps))  # take inverse
#                 T_var = torch.matmul(T_var, marginals)
#             else:
#                 # marginals_alpha = T_var @ torch.ones(T_var.shape[1], dtype=T_var.dtype).to(device)
#                 marginals_beta = T_var.t() @ torch.ones(T_var.shape[0], dtype=T_var.dtype).to(device)
#
#                 marginals = (1 / (marginals_beta + eps))
# #                 print("shape of inverse marginals beta is ", marginals_beta.shape)
# #                 print("inverse marginals beta is ", marginals_beta)
#
#                 T_var = T_var * marginals
#                 # i.e., how a neuron of 2nd model is constituted by the neurons of 1st model
#                 # this should all be ones, and number equal to number of neurons in 2nd model
#                 print(T_var.sum(dim=0))
#                 # assert (T_var.sum(dim=0) == torch.ones(T_var.shape[1], dtype=T_var.dtype).to(device)).all()
#
#         if args.debug:
#             if idx == (num_layers - 1):
#                 print("there goes the last transport map: \n ", T_var)
#             else:
#                 print("there goes the transport map at layer {}: \n ".format(idx), T_var)
#
#             print("Ratio of trace to the matrix sum: ", torch.trace(T_var) / torch.sum(T_var))
#
# #         print("Ratio of trace to the matrix sum: ", torch.trace(T_var) / torch.sum(T_var))
# #         print("Here, trace is {} and matrix sum is {} ".format(torch.trace(T_var), torch.sum(T_var)))
#         setattr(args, 'trace_sum_ratio_{}'.format(layer0_name), (torch.trace(T_var) / torch.sum(T_var)).item())
#
#         if args.past_correction:
# #             print("this is past correction for weight mode")
# #             print("Shape of aligned wt is ", aligned_wt.shape)
# #             print("Shape of fc_layer0_weight_data is ", fc_layer0_weight_data.shape)
#             t_fc0_model = torch.matmul(T_var.t(), aligned_wt.contiguous().view(aligned_wt.shape[0], -1))
#             if pp1:
#                 pp1[idx+1] = torch.mm(pp1[idx+1],T_var)
# #             print("******",str(idx),"**********",T_var.shape)
#         else:
#             t_fc0_model = torch.matmul(T_var.t(), fc_layer0_weight_data.view(fc_layer0_weight_data.shape[0], -1))
#             if pp1:
#                 pp1[idx+1] = torch.mm(pp1[idx+1],T_var)
#
#         # Average the weights of aligned first layers
#
#         if args.ensemble_step != 0.5:
#             geometric_fc = t_fc0_model
# #             geometric_fc = ((1-args.ensemble_step) * t_fc0_model +
# #                             args.ensemble_step * fc_layer1_weight_data.view(fc_layer1_weight_data.shape[0], -1))
#         else:
#             geometric_fc = t_fc0_model
# #             geometric_fc = (t_fc0_model + fc_layer1_weight_data.view(fc_layer1_weight_data.shape[0], -1))/2
#         if is_conv and layer_shape != geometric_fc.shape:
#             geometric_fc = geometric_fc.view(layer_shape)
#         avg_aligned_layers.append(geometric_fc)
#
#         # get the performance of the model 0 aligned with respect to the model 1
# #         if args.eval_aligned:
# #             if is_conv and layer_shape != t_fc0_model.shape:
# #                 t_fc0_model = t_fc0_model.view(layer_shape)
# #             model0_aligned_layers.append(t_fc0_model)
# #             _, acc = update_model(args, networks[0], model0_aligned_layers, test=True,
# #                                   test_loader=test_loader, idx=0)
# #             print("For layer idx {}, accuracy of the updated model is {}".format(idx, acc))
# #             setattr(args, 'model0_aligned_acc_layer_{}'.format(str(idx)), acc)
# #             if idx == (num_layers - 1):
# #                 setattr(args, 'model0_aligned_acc', acc)
#     if pp1!=None:
#         return avg_aligned_layers,pp1
#     else:
#         return avg_aligned_layers
#
# def get_wasserstein(args, fc_layer0_weight,fc_layer1_weight, eps=1e-7, test_loader=None,pp1=None):
#
#     ground_metric_object = GroundMetric(args)
#     assert fc_layer0_weight.shape == fc_layer1_weight.shape
#
#     mu_cardinality = fc_layer0_weight.shape[0]
#     nu_cardinality = fc_layer1_weight.shape[0]
#
#
#     layer_shape = fc_layer0_weight.shape
#     if len(layer_shape) > 2:
#         is_conv = True
#         fc_layer0_weight_data = fc_layer0_weight.data.view(fc_layer0_weight.shape[0], fc_layer0_weight.shape[1], -1)
#         fc_layer1_weight_data = fc_layer1_weight.data.view(fc_layer1_weight.shape[0], fc_layer1_weight.shape[1], -1)
#     else:
#         is_conv = False
#         fc_layer0_weight_data = fc_layer0_weight.data
#         fc_layer1_weight_data = fc_layer1_weight.data
#
#     if is_conv:
#         M = ground_metric_object.process(fc_layer0_weight_data.view(fc_layer0_weight_data.shape[0], -1),
#                         fc_layer1_weight_data.view(fc_layer1_weight_data.shape[0], -1))
#     else:
#         M = ground_metric_object.process(fc_layer0_weight_data, fc_layer1_weight_data)
#
#     mu = get_histogram(args, 0, mu_cardinality)
#     nu = get_histogram(args, 1, nu_cardinality)
#     cpuM = M.data.cpu().numpy()
#     if args.exact:
#         dd = ot.emd2(mu, nu, cpuM)
#
#     else:
#         dd = ot.bregman.sinkhorn2(mu, nu, cpuM, reg=args.reg)
#
#     return dd
#
# # def print_stats(arr, nick=""):
# #     print(nick)
# #     print("summary stats are: \n max: {}, mean: {}, min: {}, median: {}, std: {} \n".format(
# #         arr.max(), arr.mean(), arr.min(), np.median(arr), arr.std()
# #     ))
#
# # def get_activation_distance_stats(activations_0, activations_1, layer_name=""):
# # #     if layer_name != "":
# # #         print("In layer {}: getting activation distance statistics".format(layer_name))
# #     M = cost_matrix(activations_0, activations_1) ** (1/2)
# #     mean_dists =  torch.mean(M, dim=-1)
# #     max_dists = torch.max(M, dim=-1)[0]
# #     min_dists = torch.min(M, dim=-1)[0]
# #     std_dists = torch.std(M, dim=-1)
#
# #     print("Statistics of the distance from neurons of layer 1 (averaged across nodes of layer 0): \n")
# #     print("Max : {}, Mean : {}, Min : {}, Std: {}".format(torch.mean(max_dists), torch.mean(mean_dists), torch.mean(min_dists), torch.mean(std_dists)))
#
# # def update_model(args, model, new_params, test=False, test_loader=None, reversed=False, idx=-1):
#
# #     updated_model = mynetwork.mnistnet()
# #     if args.gpu_id != -1:
# #         updated_model = updated_model.cuda(args.gpu_id)
#
# #     layer_idx = 0
# #     model_state_dict = model.state_dict()
#
# # #     print("len of model_state_dict is ", len(model_state_dict.items()))
# # #     print("len of new_params is ", len(new_params))
#
# #     for key, value in model_state_dict.items():
# # #         print("updated parameters for layer ", key)
# #         model_state_dict[key] = new_params[layer_idx]
# #         layer_idx += 1
# #         if layer_idx == len(new_params):
# #             break
#
#
# #     updated_model.load_state_dict(model_state_dict)
#
# # #     if test:
# # #         log_dict = {}
# # #         log_dict['test_losses'] = []
# # #         final_acc = routines.test(args, updated_model, test_loader, log_dict)
# # #         print("accuracy after update is ", final_acc)
# # #     else:
# # #          final_acc = None
#
# #     return updated_model#, final_acc
#
#
#
#
# def _get_layer_weights(layer_weight, is_conv):
#     if is_conv:
#         # For convolutional layers, it is (#out_channels, #in_channels, height, width)
#         layer_weight_data = layer_weight.data.view(layer_weight.shape[0], layer_weight.shape[1], -1)
#     else:
#         layer_weight_data = layer_weight.data
#
#     return layer_weight_data
#
#
#
#
# def _custom_sinkhorn(args, mu, nu, cpuM):
#     if not args.unbalanced:
#         if args.sinkhorn_type == 'normal':
#             T = ot.bregman.sinkhorn(mu, nu, cpuM, reg=args.reg)
#         elif args.sinkhorn_type == 'stabilized':
#             T = ot.bregman.sinkhorn_stabilized(mu, nu, cpuM, reg=args.reg)
#         elif args.sinkhorn_type == 'epsilon':
#             T = ot.bregman.sinkhorn_epsilon_scaling(mu, nu, cpuM, reg=args.reg)
# #         elif args.sinkhorn_type == 'gpu':
# #             T, _ = utils.sinkhorn_loss(cpuM, mu, nu, gpu_id=args.gpu_id, epsilon=args.reg, return_tmap=True)
#         else:
#             raise NotImplementedError
#     else:
#         T = ot.unbalanced.sinkhorn_knopp_unbalanced(mu, nu, cpuM, reg=args.reg, reg_m=args.reg_m)
#     return T
#
#
# def _sanity_check_tmap(T):
#     if not math.isclose(np.sum(T), 1.0, abs_tol=1e-7):
# #         print("Sum of transport map is ", np.sum(T))
#         raise Exception('NAN inside Transport MAP. Most likely due to large ground metric values')
#
#
#
# # def _check_layer_sizes(args, layer_idx, shape1, shape2, num_layers):
# #     if args.width_ratio == 1:
# #         return shape1 == shape2
# #     else:
# #         if args.dataset == 'mnist':
# #             if layer_idx == 0:
# #                 return shape1[-1] == shape2[-1] and (shape1[0]/shape2[0]) == args.width_ratio
# #             elif layer_idx == (num_layers -1):
# #                 return (shape1[-1]/shape2[-1]) == args.width_ratio and shape1[0] == shape2[0]
# #             else:
# #                 ans = True
# #                 for ix in range(len(shape1)):
# #                     ans = ans and shape1[ix]/shape2[ix] == args.width_ratio
# #                 return ans
# #         elif args.dataset[0:7] == 'Cifar10':
# #             assert args.second_model_name is not None
# #             if layer_idx == 0 or layer_idx == (num_layers -1):
# #                 return shape1 == shape2
# #             else:
# #                 if (not args.reverse and layer_idx == (num_layers-2)) or (args.reverse and layer_idx == 1):
# #                     return (shape1[1] / shape2[1]) == args.width_ratio
# #                 else:
# #                     return (shape1[0]/shape2[0]) == args.width_ratio
#
#
# def _compute_marginals(args, T_var, device, eps=1e-7):
#     if args.correction:
#         if not args.proper_marginals:
#             # think of it as m x 1, scaling weights for m linear combinations of points in X
#             marginals = torch.ones(T_var.shape)
#             if args.gpu_id != -1:
#                 marginals = marginals.cuda(args.gpu_id)
#
#             marginals = torch.matmul(T_var, marginals)
#             marginals = 1 / (marginals + eps)
# #             print("marginals are ", marginals)
#
#             T_var = T_var * marginals
#
#         else:
#             # marginals_alpha = T_var @ torch.ones(T_var.shape[1], dtype=T_var.dtype).to(device)
#             marginals_beta = T_var.t() @ torch.ones(T_var.shape[0], dtype=T_var.dtype).to(device)
#
#             marginals = (1 / (marginals_beta + eps))
# #             print("shape of inverse marginals beta is ", marginals_beta.shape)
# #             print("inverse marginals beta is ", marginals_beta)
#
#             T_var = T_var * marginals
#             # i.e., how a neuron of 2nd model is constituted by the neurons of 1st model
#             # this should all be ones, and number equal to number of neurons in 2nd model
#             print(T_var.sum(dim=0))
#             # assert (T_var.sum(dim=0) == torch.ones(T_var.shape[1], dtype=T_var.dtype).to(device)).all()
#
# #         print("T_var after correction ", T_var)
# #         print("T_var stats: max {}, min {}, mean {}, std {} ".format(T_var.max(), T_var.min(), T_var.mean(),
# #                                                                      T_var.std()))
#     else:
#         marginals = None
#
#     return T_var, marginals
#
# def _get_current_layer_transport_map(args, mu, nu, M0, M1, idx, layer_shape, eps=1e-7, layer_name=None):
#
#     if not args.gromov:
#         cpuM = M0.data.cpu().numpy()
#         if args.exact:
#             T = ot.emd(mu, nu, cpuM)
#         else:
#             T = _custom_sinkhorn(args, mu, nu, cpuM)
#
#         if args.print_distances:
#             ot_cost = np.multiply(T, cpuM).sum()
# #             print(f'At layer idx {idx} and shape {layer_shape}, the OT cost is ', ot_cost)
#             if layer_name is not None:
#                 setattr(args, f'{layer_name}_layer_{idx}_cost', ot_cost)
#             else:
#                 setattr(args, f'layer_{idx}_cost', ot_cost)
#     else:
#         cpuM0 = M0.data.cpu().numpy()
#         cpuM1 = M1.data.cpu().numpy()
#
#         assert not args.exact
#         T = ot.gromov.entropic_gromov_wasserstein(cpuM0, cpuM1, mu, nu, loss_fun=args.gromov_loss, epsilon=args.reg)
#
#     if not args.unbalanced:
#         _sanity_check_tmap(T)
#
#     if args.gpu_id != -1:
#         T_var = torch.from_numpy(T).cuda(args.gpu_id).float()
#     else:
#         T_var = torch.from_numpy(T).float()
#
# #     if args.tmap_stats:
# #         print(
# #         "Tmap stats (before correction) \n: For layer {}, frobenius norm from the joe's transport map is {}".format(
# #             layer0_name, torch.norm(T_var - torch.ones_like(T_var) / torch.numel(T_var), p='fro')
# #         ))
#
# #     print("shape of T_var is ", T_var.shape)
# #     print("T_var before correction ", T_var)
#
#     return T_var
#
# def _get_neuron_importance_histogram(args, layer_weight, is_conv, eps=1e-9):
# #     print('shape of layer_weight is ', layer_weight.shape)
#     if is_conv:
#         layer = layer_weight.contiguous().view(layer_weight.shape[0], -1).cpu().numpy()
#     else:
#         layer = layer_weight.cpu().numpy()
#
#     if args.importance == 'l1':
#         importance_hist = np.linalg.norm(layer, ord=1, axis=-1).astype(
#                     np.float64) + eps
#     elif args.importance == 'l2':
#         importance_hist = np.linalg.norm(layer, ord=2, axis=-1).astype(
#                     np.float64) + eps
#     else:
#         raise NotImplementedError
#
#     if not args.unbalanced:
#         importance_hist = (importance_hist/importance_hist.sum())
# #         print('sum of importance hist is ', importance_hist.sum())
#     # assert importance_hist.sum() == 1.0
#     return importance_hist
#
#
# def get_network_from_param_list(args, param_list, mtype,test_loader):
#
# #     print("using independent method")
#     if mtype=='mlpnet':
#         new_network = mynetwork.mnistnet()#get_model_from_name(args, idx=1)
#     elif mtype=='cnnnet':
#         new_network = mynetwork.cnnNet()#get_model_from_name(args, idx=1)
#     if args.gpu_id != -1:
#         new_network = new_network.cuda(args.gpu_id)
#
#     # check the test performance of the network before
# #     log_dict = {}
# #     log_dict['test_losses'] = []
# #     routines.test(args, new_network, test_loader, log_dict)
#
#     # set the weights of the new network
#     # print("before", new_network.state_dict())
# #     print("len of model parameters and avg aligned layers is ", len(list(new_network.parameters())),
# #           len(param_list))
#     assert len(list(new_network.parameters())) == len(param_list)
#
#     layer_idx = 0
#     model_state_dict = new_network.state_dict()
#
# #     print("len of model_state_dict is ", len(model_state_dict.items()))
# #     print("len of param_list is ", len(param_list))
#
#     for key, value in model_state_dict.items():
#         model_state_dict[key] = param_list[layer_idx]
#         layer_idx += 1
#
#     new_network.load_state_dict(model_state_dict)
#
#     # check the test performance of the network after
# #     log_dict = {}
# #     log_dict['test_losses'] = []
# #     acc = routines.test(args, new_network, test_loader, log_dict)
#
#     return  new_network
#
# def geometric_ensembling_modularized(args, networks, train_loader=None,mtype='mlpnet' ,test_loader=None,pp1=None):
#
#     if pp1!=None:
#         avg_aligned_layers,pp1 = get_wassersteinized_layers_modularized(args, networks, test_loader=test_loader,pp1=pp1)
#         mmm = get_network_from_param_list(args, avg_aligned_layers,mtype, test_loader)
#         return mmm,pp1
#     else:
#         avg_aligned_layers = get_wassersteinized_layers_modularized(args, networks, test_loader=test_loader)
#         return get_network_from_param_list(args, avg_aligned_layers,mtype, test_loader)
#
#