def loss_lookahead_diff(model: NeuralTeleportationModel, data: Tensor, target: Tensor, metrics: TrainingMetrics, config: OptimalTeleportationTrainingConfig, **kwargs) -> Number: # Save the state of the model, prior to performing the lookahead state_dict = model.state_dict() # Initialize a new optimizer to perform lookahead optimizer = get_optimizer_from_model_and_config(model, config) optimizer.zero_grad() # Compute loss at the teleported point loss = torch.stack([metrics.criterion(model(data_batch), target_batch) for data_batch, target_batch in zip(data, target)]).mean(dim=0) # Take a step using the gradient at the teleported point loss.backward() # Compute loss after the optimizer step lookahead_loss = torch.stack([metrics.criterion(model(data_batch), target_batch) for data_batch, target_batch in zip(data, target)]).mean(dim=0) # Restore the state of the model prior to the lookahead model.load_state_dict(state_dict) # Compute the difference between the lookahead loss and the original loss return (loss - lookahead_loss).item()
def teleport_model_to_optimize_metric(model: NeuralTeleportationModel, train_dataset: Dataset, metrics: TrainingMetrics, config: "OptimalTeleportationTrainingConfig", **kwargs) \ -> NeuralTeleportationModel: print(f"Selecting best of {config.num_teleportations} random COBs " f"w.r.t. {config.optim_metric.__name__}") # Extract a single batch on which to compute gradients for each model to be compared dataloader = DataLoader(train_dataset, batch_size=config.batch_size) data, target = [], [] for (data_batch, target_batch), _ in zip(dataloader, range(config.num_batches)): data.append(data_batch) target.append(target_batch) data = torch.stack(data).to(device=config.device) target = torch.stack(target).to(device=config.device) optimal_metric = config.optim_metric(model=model, data=data, target=target, metrics=metrics, config=config) model.cpu() # Move model to CPU to avoid having 2 models on the GPU (to avoid possible CUDA OOM error) optimal_model = model for _ in range(config.num_teleportations): teleported_model = deepcopy(model).random_teleport(cob_range=config.cob_range, sampling_type=config.cob_sampling) teleported_model.to(config.device) # Move model back to chosen device before computing gradients metric = config.optim_metric(model=teleported_model, data=data, target=target, metrics=metrics, config=config) teleported_model.cpu() # Move model back to CPU after computation is done (to avoid possible CUDA OOM error) if metric > optimal_metric: optimal_model = teleported_model optimal_metric = metric return optimal_model.to(config.device)
def generate_teleportation_training_weights( model: NeuralTeleportationModel, trainset: Dataset, metric: TrainingMetrics, config: LandscapeConfig) -> Tuple[List[torch.Tensor], torch.Tensor]: """ This will generate a list of weights at a given epoch while training the passed model. If teleport_every is different than 0, the model will teleport every time. """ w = [model.get_weights().clone().detach().cpu()] trainloader = torch.utils.data.DataLoader(trainset, batch_size=config.batch_size, drop_last=True) optim = get_optimizer_from_model_and_config(model, config) for e in range(config.epochs): if e in config.teleport_at and config.teleport_at != 0: print("Teleporting Model...") model.random_teleport(cob_range=config.cob_range, sampling_type=config.cob_sampling) w.append(model.get_weights().clone().detach().cpu()) optim = get_optimizer_from_model_and_config(model, config) train_epoch(model, metrics=metric, config=config, train_loader=trainloader, optimizer=optim, epoch=e, device=config.device) w.append(model.get_weights().clone().detach().cpu()) final = w[-1::][0] return w, final
def test_model_without_set_get_weights( model: nn.Module, testset: Dataset, metric: TrainingMetrics, config: TrainingConfig, rept: int = 1) -> Tuple[np.ndarray, np.ndarray]: """ Test if the model is coequal before and after using model.teleport """ loss_diff_avg = [] acc_diff_avg = [] for _ in range(rept): m = NeuralTeleportationModel(model, input_shape=(config.batch_size, 3, 32, 32)).to(device) res = test(m, testset, metric, config) loss1, acc1 = res['loss'], res['accuracy'] m.random_teleport() res = test(m, testset, metric, config) loss2, acc2 = res['loss'], res['accuracy'] loss_diff_avg.append(np.abs(loss1 - loss2)) acc_diff_avg.append(np.abs(acc1 - acc2)) print("==========================================") print("Loss and accuracy diff without set/get was") print("Loss diff was: {:.6e}".format(np.abs(loss1 - loss2))) print("Acc diff was: {:.6e}".format(np.abs(acc1 - acc2))) print("==========================================") return np.mean(loss_diff_avg), np.mean(acc_diff_avg)
def test_model_with_set_get_weights( model: nn.Module, testset: Dataset, metric: TrainingMetrics, config: TrainingConfig, rept: int = 1) -> Tuple[np.ndarray, np.ndarray]: loss_diff_avg = [] acc_diff_avg = [] for _ in range(rept): m = NeuralTeleportationModel(model, input_shape=(config.batch_size, 3, 32, 32)).to(device) w_o, cob_o = m.get_params() m.random_teleport() w_t, cob_t = m.get_params() m.set_params(weights=w_o, cob=cob_o) res = test(m, testset, metric, config) loss1, acc1 = res['loss'], res['accuracy'] m.set_params(weights=w_t, cob=cob_t) res = test(m, testset, metric, config) loss2, acc2 = res['loss'], res['accuracy'] loss_diff_avg.append(np.abs(loss1 - loss2)) acc_diff_avg.append(np.abs(acc1 - acc2)) print("==========================================") print("Loss and accuracy diff with set/get was") print("Loss diff was: {:.6e}".format(np.abs(loss1 - loss2))) print("Acc diff was: {:.6e}".format(np.abs(acc1 - acc2))) print("==========================================") return np.mean(loss_diff_avg), np.mean(acc_diff_avg)
def generate_1D_linear_interp( model: NeuralTeleportationModel, param_o: Tuple[torch.Tensor, torch.Tensor], param_t: Tuple[torch.Tensor, torch.Tensor], a: torch.Tensor, trainset: Dataset, valset: Dataset, metric: TrainingMetrics, config: TrainingConfig, checkpoint: dict = None) -> Tuple[list, list, list, list]: """ This is 1-Dimensional Linear Interpolation θ(α) = (1−α)θ + αθ′ """ loss = [] loss_v = [] acc_t = [] acc_v = [] w_o, cob_o = param_o w_t, cob_t = param_t start_at = checkpoint["step"] if checkpoint else 0 try: for step, coord in enumerate(a, start_at): # Interpolate the weight from W to T(W), # then interpolate the cob for the activation # and batchNorm layers only. print("step {} of {} - alpha={}".format(step + 1, len(a), coord)) w = (1 - coord) * w_o + coord * w_t cob = (1 - coord) * cob_o + coord * cob_t model.set_params(w, cob) res = test(model, trainset, metric, config) loss.append(res['loss']) acc_t.append(res['accuracy']) res = test(model, valset, metric, config) acc_v.append(res['accuracy']) loss_v.append(res['loss']) except: if not checkpoint: checkpoint = { 'step': step, 'alpha': a, 'original_model': param_o, 'teleported_model': param_t, 'losses': loss, 'acc_t': acc_t, 'acc_v': acc_v, } else: checkpoint['step'] = step checkpoint['losses'] = checkpoint['losses'].append(loss) checkpoint['acc_t'] = checkpoint['acc_t'].append(acc_t) checkpoint['acc_v'] = checkpoint['acc_v'].append(loss) torch.save(checkpoint, linterp_checkpoint_file) print("A checkpoint was made on step {} of {}".format(step, len(a))) # This is to notify the upper level of try/except # Since there is no way to know if this is from before teleportation or after teleportation. raise return loss, acc_t, loss_v, acc_v
def teleport_model_randomly(model: NeuralTeleportationModel, config: "RandomTeleportationTrainingConfig", **kwargs) \ -> NeuralTeleportationModel: if random.random() < config.teleport_prob: print("Applying random COB to model in training") model.random_teleport(cob_range=config.cob_range, sampling_type=config.cob_sampling) else: print("Skipping COB") return model
def simulate_teleportation_distribution( model: NeuralTeleportationModel, config: "DistributionTeleportationTrainingConfig", **kwargs) -> NeuralTeleportationModel: print( f"Shifting weights to a similar distribution to a {config.cob_sampling} teleportation w/ {config.cob_range}" f"COB range.") model.cpu() teleported_model = deepcopy(model).random_teleport( cob_range=config.cob_range, sampling_type=config.cob_sampling) teleported_weights = teleported_model.get_weights( concat=True).cpu().detach().numpy() hist, bin_edges = np.histogram(teleported_weights, bins=1000) hist = hist / hist.sum() model.init_like_histogram(hist, bin_edges) return model.to(config.device)
def weighted_grad_norm(model: NeuralTeleportationModel, data: Tensor, target: Tensor, metrics: TrainingMetrics, order: Union[str, number] = 'fro', **kwargs) -> Number: weights = model.get_weights() gradients = torch.stack([model.get_grad(data_batch, target_batch, metrics.criterion) for data_batch, target_batch in zip(data, target)]).mean(dim=0) # Compute the gradient/weight ratio where possible ratio = gradients / weights # Identify where the ratio is numerically unstable (division by 0-valued weights) nan_ratio_mask = torch.isnan(ratio) # Replace unstable values by statistically representative measures ratio[nan_ratio_mask] = ratio[~nan_ratio_mask].mean() # Compute the norm of the ratio and move result to CPU (to avoid cluttering GPU if fct is called repeatedly) return torch.norm(ratio, p=order).item()
def generate_contour_loss_values( model: NeuralTeleportationModel, directions: Tuple[torch.Tensor, torch.Tensor], weights: torch.Tensor, surface: torch.Tensor, trainset: Dataset, metric: TrainingMetrics, config: TrainingConfig, checkpoint: dict = None) -> Tuple[np.ndarray, np.ndarray]: """ Generate a tensor containing the loss values from a given model. """ loss = [] acc = [] delta, eta = directions start_at = 0 if checkpoint: start_at = checkpoint['step'] try: for step, (x, y) in enumerate(surface, start_at): print("Evaluating step {}: [{:.3f}, {:.3f}]".format(step, x, y)) x, y = x.to(config.device), y.to(config.device) # L (w + alpha*delta + beta*eta) changes = (delta * x + eta * y).to(config.device) w = torch.add(weights, changes) model.set_weights(w) results = test(model, trainset, metric, config) loss.append(results['loss']) acc.append(results['accuracy']) except: # The reason is that, no matter what, make a checkpoint of the current surface generation. if not checkpoint: checkpoint = {'step': step, 'surface': surface, 'loss': loss} else: checkpoint['step'] = step [checkpoint['loss'].append(l) for l in loss] torch.save(checkpoint, contour_checkpoint_file) print("A checkpoint was made at coord {} of {}".format(x, y)) # This is to notify the upper level of try/except # Since there is no way to know if this is from before teleportation or after teleportation. raise return np.array(loss), np.array(acc)
def get_model(dataset_name: str, model_name: str, device: str = 'cpu', initializer: Dict[str, Union[str, float]] = None, **model_kwargs) -> NeuralTeleportationModel: # Look up if the requested model is available in the model zoo model_factories = _get_model_factories() if model_name not in model_factories: raise KeyError(f"{model_name} was not found in the model zoo") # Dynamically determine the parameters for initializing the model based on the dataset model_kwargs.update(get_dataset_info(dataset_name, "num_classes")) if "mlp" in model_name.lower(): input_channels, image_size = get_dataset_info(dataset_name, "input_channels", "image_size").values() model_kwargs.update(input_shape=(input_channels, *image_size)) else: model_kwargs.update(get_dataset_info(dataset_name, "input_channels")) if "cifar" in dataset_name and ("resnet" in model_name): model_kwargs.update({"for_dataset": "cifar"}) # Instantiate the model model_factory = model_factories[model_name] model = model_factory(**model_kwargs) # Initialize the model if initializer is not None: init_gain = None if "gain" not in initializer.keys( ) and initializer["type"] == "none" else initializer["gain"] init_non_linearity = None if "non_linearity" not in initializer.keys( ) else initializer["non_linearity"] model = initialize_model(model, init_type=initializer["type"], init_gain=init_gain, non_linearity=init_non_linearity) # Transform the base ``nn.Module`` to a ``NeuralTeleportationModel`` input_channels, image_size = get_dataset_info(dataset_name, "input_channels", "image_size").values() model = NeuralTeleportationModel(network=model, input_shape=(2, input_channels, *image_size)) return model.to(device)
def simulate_teleportation_sphere(model: NeuralTeleportationModel, config: "PseudoTeleportationTrainingConfig", **kwargs) -> NeuralTeleportationModel: print( f"Shifting weights on a sphere similar to a {config.cob_sampling} teleportation w/ {config.cob_range} " f"COB range.") model.cpu( ) # Move model to CPU to avoid having 2 models on the GPU (to avoid possible CUDA OOM error) teleported_model = deepcopy(model).random_teleport( cob_range=config.cob_range, sampling_type=config.cob_sampling) init_layers = model.get_weights(concat=False) teleported_layers = teleported_model.get_weights(concat=False) pseudo_teleported_layers = [] for init_layer, teleported_layer in zip(init_layers, teleported_layers): layer_shift = torch.randn_like(init_layer) layer_shift = normalize(layer_shift, p=1, dim=0) * torch.norm( teleported_layer - init_layer, 1) pseudo_teleported_layer = init_layer + layer_shift pseudo_teleported_layers.append(pseudo_teleported_layer) pseudo_teleported_weights = torch.cat(pseudo_teleported_layers) model.set_weights(pseudo_teleported_weights) return model.to(config.device)
def test_set_weights(network: nn.Module, input_shape: Tuple = (1, 1, 28, 28), model_name: str = None): """ test_set_weights checks if method set_weights() in NeuralTeleportationModel works Args: network (nn.Module): Network to test input_shape (tuple): Input shape of network model_name (str): The name or label assigned to differentiate the model """ model_name = model_name or network.__class__.__name__ model = NeuralTeleportationModel(network, input_shape) w1 = model.get_weights() model.reset_weights() model.set_weights(w1) w2 = model.get_weights() assert np.allclose(w1.detach().numpy(), w2.detach().numpy()) print("Weights set successfully for " + model_name + " model.")
def test_teleport(network: nn.Module, input_shape: Tuple = (1, 1, 28, 28), verbose: bool = False, atol: float = 1e-5, model_name: str = None): """ Return mean of the difference between the weights of network and a random teleportation, and checks if teleportation has the same network function Args: network (nn.Module): Network to be tested input_shape (tuple): Input shape of network verbose (bool): Flag to print comparision between network and a teleportation atol (float): Absolute tolerance allowed between outputs to pass the test model_name (str): The name or label assigned to differentiate the model Returns: float with the average of the difference between the weights of the network and a teleportation """ model_name = model_name or network.__class__.__name__ model = NeuralTeleportationModel(network=network, input_shape=input_shape) model.eval() # model must be set to eval because of dropout x = torch.rand(input_shape) pred1 = model(x).detach().numpy() w1 = model.get_weights().detach().numpy() model.random_teleport() pred2 = model(x).detach().numpy() w2 = model.get_weights().detach().numpy() diff_average = np.mean(np.abs((pred1 - pred2))) if verbose: print("Sample outputs: ") print("Pre teleportation: ", pred1.flatten()[:10]) print("Post teleportation: ", pred2.flatten()[:10]) print("Diff weight average: ", np.mean(np.abs((w1 - w2)))) print("Diff prediction average: ", diff_average) assert not np.allclose(w1, w2) assert np.allclose(pred1, pred2, atol=atol), "Teleporation did not work for model {}. Average difference: {}". \ format(model_name, diff_average) print("Teleportation successful for " + model_name + " model.") return diff_average
def test_calculate_ones(network, model_name=None, input_shape=(1, 1, 28, 28), noise=False, verbose=False): """ Test if the correct change of basis can be calculated for a cob of ones. Args: network (nn.Module): Network to be tested model_name (str): The name or label assigned to differentiate the model input_shape (tuple): Input shape of network noise (bool): whether to add noise to the target weights before optimisation. verbose (bool): whether to display sample ouputs during the test """ model_name = model_name or network.__class__.__name__ model = NeuralTeleportationModel(network=network, input_shape=input_shape) model.initialize_cob() w1 = model.get_weights(concat=False, flatten=False, bias=False) _w1 = model.get_weights(concat=False, flatten=False, bias=False) if noise: for w in _w1: w += torch.rand(w.shape) * 0.001 cob = model.get_cob() calculated_cob = model.calculate_cob(w1, _w1) error = (cob - calculated_cob).abs().mean() if verbose: print("Cob: ", cob.flatten()[:10]) print("Calculated cob: ", calculated_cob.flatten()[:10]) print("cob error ", (calculated_cob - cob).flatten()[:10]) print("cob error : ", error) assert np.allclose( cob, calculated_cob ), "Calculate cob (ones) FAILED for " + model_name + " model." print("Calculate cob (ones) successful for " + model_name + " model.")
def test_multiple_teleport(network: nn.Module, input_shape: Tuple = (1, 1, 28, 28), verbose: bool = False, atol: float = 1e-5, model_name: str = None): """ Test multiple successive teleporations. Args: network (nn.Module): Network to be tested input_shape (tuple): Input shape of network verbose (bool): Flag to print comparision between network and a teleportation atol (float): Absolute tolerance allowed between outputs to pass the test model_name (str): The name or label assigned to differentiate the model """ model_name = model_name or network.__class__.__name__ model = NeuralTeleportationModel(network=network, input_shape=input_shape) model.eval() # model must be set to eval because of dropout x = torch.rand(input_shape) pred1 = model(x).detach().numpy() for _ in range(10): model.random_teleport(cob_range=10, sampling_type='inter_landscape') pred2 = model(x).detach().numpy() diff_average = np.mean(np.abs((pred1 - pred2))) assert np.allclose(pred1, pred2, atol=atol), "Multiple Teleporation did not work for model {}. Average difference: {}".format( model_name, diff_average) for _ in range(10): model.random_teleport(cob_range=10, sampling_type='inter_landscape', reset_teleportation=False) pred2 = model(x).detach().numpy() diff_average = np.mean(np.abs((pred1 - pred2))) assert np.allclose(pred1, pred2, atol=atol), "Multiple Teleporation did not work for model {}. Average difference: {}".format( model_name, diff_average) print("Multiple Teleportations successful for " + model_name + " model.")
def start_training(model: NeuralTeleportationModel, trainloader: DataLoader, valset: VisionDataset, metric: TrainingMetrics, config: CompareTrainingConfig, teleport_chance: float) -> np.ndarray: """ This function starts a model training with a specific Scenario configuration. Scenario 1: train the model without using teleportation (teleportation_chance = 0.0) Scenario 2: train the model using a probability of teleporting every Xth epochs (0 < teleportation_chance < 1.0) Scenario 3: train the model using teleportation every Xth epochs (teleportation_chance = 1.0) returns: np.array containing the validation accuracy results of every epochs. """ model.to(config.device) optimizer = get_optimizer_from_model_and_config(model, config) results = [] for e in np.arange(1, args.epochs + 1): train_epoch(model=model, metrics=metric, optimizer=optimizer, train_loader=trainloader, epoch=e, device=config.device) results.append(test(model=model, dataset=valset, metrics=metric, config=config)['accuracy']) model.train() if e % config.every_n_epochs == 0 and random.random() <= teleport_chance: print("teleported model") if config.targeted_teleportation: # TODO: use teleportation function here when they are available. raise NotImplementedError else: model.random_teleport(cob_range=config.cob_range, sampling_type=config.cob_sampling) optimizer = get_optimizer_from_model_and_config(model, config) model.cpu() # Force the network to go out of the cuda mem. return np.array(results)
def test_cuda_teleport(network, input_shape=(1, 1, 28, 28), verbose=False): """ Test if a model can be teleported successfully on cuda. Args: network (nn.Module): Model to test input_shape (tuple): Input shape for the model verbose (bool): if True samples of predictions are printed Returns: Average difference between elements of prediction before and after teleportation. """ network = network.cuda() model = NeuralTeleportationModel(network=network, input_shape=input_shape) x = torch.rand(input_shape).cuda() pred1 = model(x).cpu().detach().numpy() w1 = model.get_weights().cpu().detach().numpy() model.random_teleport() pred2 = model(x).cpu().detach().numpy() w2 = model.get_weights().cpu().detach().numpy() diff_average = (w1 - w2).mean() if verbose: print("Model on device: {}".format(next(network.parameters()).device)) print("Sample outputs: ") print("Pre teleportation: ", pred1.flatten()[:10]) print("Post teleportation: ", pred2.flatten()[:10]) assert not np.allclose(w1, w2) assert np.allclose( pred1, pred2), "Teleporation did not work. Average difference: {}".format( diff_average) print("Teleportation successful.") return diff_average
from neuralteleportation.models.model_zoo.vggcob import vgg16_bnCOB from neuralteleportation.models.model_zoo.resnetcob import resnet18COB from torch.utils.data import DataLoader from tqdm import tqdm device = 'cuda' if torch.cuda.is_available() else 'cpu' input_shape = (32, 3, 32, 32) trainset, valset, testset = experiment_setup.get_dataset_subsets("cifar10") train_loader = DataLoader(trainset, batch_size=input_shape[0], shuffle=True) # MLP pbar = tqdm(enumerate(train_loader)) mlp = MLPCOB(input_shape=(3, 32, 32), num_classes=10).to(device=device) mlp = NeuralTeleportationModel(network=mlp, input_shape=input_shape) plot_difference_teleported_gradients(network=mlp, pbar=pbar, network_descriptor='MLP on CIFAR-10', device=device) pbar.close() # VGG pbar = tqdm(enumerate(train_loader)) vgg = vgg16_bnCOB(num_classes=10).to(device=device) vgg = NeuralTeleportationModel(network=vgg, input_shape=input_shape) plot_difference_teleported_gradients(network=vgg, pbar=pbar, network_descriptor='VGG on CIFAR-10', device=device) pbar.close() # ResNet pbar = tqdm(enumerate(train_loader)) resnet = resnet18COB(num_classes=10).to(device=device) resnet = NeuralTeleportationModel(network=resnet, input_shape=input_shape) plot_difference_teleported_gradients(network=resnet, pbar=pbar, network_descriptor='ResNet on CIFAR-10',
def plot_model_weights_histogram(model: NeuralTeleportationModel, mode: str, title: str, output_dir: Path = None, save_format: str = None, xlim: float = None, ylim_max: float = None, zoom_plot: bool = False) -> None: def _format_ticklabels(ticklabels) -> List[str]: return [f"{ticklabel:.1f}" for ticklabel in ticklabels] def _zoom_plot(ax, data, inset, lims): with sns.axes_style({'axes.linewidth': 1, 'axes.edgecolor': 'black'}): axins = ax.inset_axes(inset) sns.kdeplot(data, fill=True, ax=axins) axins.set_xlim(lims[:2]) axins.set_ylim(lims[2:]) axins.set(xticklabels=[], ylabel=None, yticklabels=[]) rectangle_patch, connector_lines = ax.indicate_inset_zoom(axins) # Make the indicator and connectors more easily visible rectangle_patch.set_linewidth(2) for connector_line in connector_lines: connector_line.set_linewidth(2) # Manually set the visibility of the appropriate connector lines connector_lines[0].set_visible(True) # Lower left connector_lines[1].set_visible(False) # Upper left connector_lines[2].set_visible(True) # Lower right connector_lines[3].set_visible(False) # Upper right return axins def _plot_1d_array_histogram(array: np.ndarray, title: str): with sns.axes_style("darkgrid"): axes = sns.kdeplot(array, fill=True) if xlim: axes.set(xlim=(-xlim, xlim)) if ylim_max: axes.set(ylim=(0, ylim_max)) axes.set(ylabel=None, yticklabels=[]) axes.set_xticklabels(_format_ticklabels(axes.get_xticks()), size=20) if zoom_plot: _zoom_plot(axes, array, inset=[0.05, 0.6, 0.35, 0.35], lims=[-0.13, -0.12, 0, 0.5]) _zoom_plot(axes, array, inset=[0.6, 0.6, 0.35, 0.35], lims=[0.12, 0.13, 0, 0.5]) if save_format: plt.savefig(output_dir / f"{title}.{save_format}", bbox_inches='tight') plt.close() else: plt.show() print(f"Plotting {title} histogram ...") if mode == "modelwise": _plot_1d_array_histogram(model.get_weights().cpu().detach().numpy(), title) elif mode == "layerwise": # Get weights of each layer, without the bias weights layers = model.get_weights(concat=False, bias=False) # Get rid of layers connected to input and output nodes, to keep only hidden layers layers = layers[1:-1] for idx, layer_weights in enumerate(layers): _plot_1d_array_histogram(layer_weights.cpu().detach().numpy(), title + f"_layer{idx}") else: raise ValueError( f"Mode {mode} is not a valid option. Choose one of: {{modelwise,layerwise}}" )
cob_range=args.cob_range, cob_sampling=args.cob_sampling, teleport_at=args.teleport_at, device=device) trainset, valset, testset = get_dataset_subsets("cifar10") model = get_model("cifar10", args.model) x = torch.linspace(-1, 1, args.x) y = torch.linspace(-1, 1, args.y) shape = x.shape if y is None else (len(x), len(y)) surface = torch.stack((x, y)) model.to(device) model = NeuralTeleportationModel(model, input_shape=(config.batch_size, 3, 32, 32)).to(device) original_w = model.get_weights() w_checkpoints, final_w = losslandscape.generate_teleportation_training_weights( model, trainset, metric=metric, config=config) delta = losslandscape.generate_random_2d_vector(final_w, seed=1) eta = losslandscape.generate_random_2d_vector(final_w, seed=2) # Calculate angle between the two direction vectors. print("angle between direction is {} rad".format( losslandscape.compute_angle(delta, eta))) loss, _ = losslandscape.generate_contour_loss_values( model, (delta, eta), surface, trainset, metric, config) loss = np.array(loss) loss = np.resize(loss, shape)
parser.add_argument( "--cob_range", type=float, default=1, help='Range for the teleportation to create target weights') return parser.parse_args() if __name__ == '__main__': args = argument_parser() torch.manual_seed(args.seed) model = NeuralTeleportationModel(network=MLPCOB(input_shape=(1, 28, 28), num_classes=10), input_shape=(1, 1, 28, 28)) # Get the initial set of weights and teleport. initial_weights = model.get_weights() model.random_teleport(cob_range=args.cob_range) # Get second set of weights (target weights) target_weights = model.get_weights() # Get the change of basis that created this set of weights. target_cob = model.get_cob(concat=True) # Generate a new random cob cob = model.generate_random_cob(cob_range=args.cob_range, requires_grad=True)
def test_set_cob(network, model_name, input_shape=(1, 1, 28, 28), verbose=False): """ Test if the set_change_of_basis method works. Args: network (nn.Module): Network to be tested model_name (str): The name or label assigned to differentiate the model input_shape (tuple): Input shape of network verbose (bool): Flag to print comparision between network and a teleportation """ x = torch.rand(input_shape) model = NeuralTeleportationModel(network, input_shape=input_shape) model.random_teleport() w1 = model.get_weights() t1 = model.get_cob() pred1 = model(x) model.reset_weights() pred2 = model(x) model.set_weights(w1) model.teleport_activations(t1) pred3 = model(x) if verbose: print("Diff prediction average: ", (pred1 - pred3).mean()) print("Pre teleportation: ", pred1.flatten()[:10]) print("Post teleportation: ", pred3.flatten()[:10]) assert not np.allclose(pred1.detach().numpy(), pred2.detach().numpy(), atol=1e-5) assert np.allclose(pred1.detach().numpy(), pred3.detach().numpy(), atol=1e-5), "Set cob/weights did not work." print("Set cob successful for " + model_name + " model.")
mnist_train, mnist_val, mnist_test = get_dataset_subsets("mnist") sample_input_shape = (1, 1, 28, 28) hidden_layers = (128, 10) net1 = MLPCOB(input_shape=(1, 28, 28), num_classes=10, hidden_layers=hidden_layers).to(device) if args.same_init: net2 = deepcopy(net1) else: net2 = MLPCOB(input_shape=(1, 28, 28), num_classes=10, hidden_layers=hidden_layers).to(device) model1 = NeuralTeleportationModel(network=net1, input_shape=sample_input_shape) if args.weights1 is not None: model1.load_state_dict(torch.load(args.weights1)) config.batch_size = 8 # Change batch size to train to different minima train(model1, train_dataset=mnist_train, metrics=metrics, config=config, val_dataset=mnist_test) torch.save(model1.state_dict(), pjoin(save_path, 'model1.pt')) print("Model 1 test results: ", test(model1, mnist_test, metrics, config)) model2 = NeuralTeleportationModel(network=net2, input_shape=sample_input_shape) if args.weights2 is not None: model2.load_state_dict(torch.load(args.weights2))
def test_calculate_cob_weights(network, model_name=None, input_shape=(1, 1, 28, 28), noise=False, verbose=True): """ Test if a cob can be calculated and applied to a network to teleport the network from the initial weights to the targets weights. Args: network (nn.Module): Network to be tested model_name (str): The name or label assigned to differentiate the model input_shape (tuple): Input shape of network noise (bool): whether to add noise to the target weights before optimisation. verbose (bool): whether to display sample ouputs during the test """ model_name = model_name or network.__class__.__name__ model = NeuralTeleportationModel(network=network, input_shape=input_shape) initial_weights = model.get_weights() w1 = model.get_weights(concat=False, flatten=False, bias=False) model.random_teleport() c1 = model.get_cob() model.random_teleport() c2 = model.get_cob() target_weights = model.get_weights() w2 = model.get_weights(concat=False, flatten=False, bias=False) if noise: for w in w2: w += torch.rand(w.shape) * 0.001 calculated_cob = model.calculate_cob(w1, w2) model.initialize_cob() model.set_weights(initial_weights) model.teleport(calculated_cob, reset_teleportation=True) calculated_weights = model.get_weights() error = (calculated_weights - initial_weights).abs().mean() if verbose: print("weights: ", target_weights.flatten()) print("Calculated cob weights: ", calculated_weights.flatten()) print("Weight error ", error) print("C1: ", c1.flatten()[:10]) print("C2: ", c2.flatten()[:10]) print("C1 * C2: ", (c1 * c2).flatten()[:10]) print("Calculated cob: ", calculated_cob.flatten()[:10]) assert np.allclose(calculated_weights.detach().numpy(), target_weights.detach().numpy()), \ "Calculate cob and weights FAILED for " + model_name + " model with error: " + str(error.item()) print("Calculate cob and weights successful for " + model_name + " model.")
def micro_teleportation_dot_product(network, dataset, nb_teleport=100, network_descriptor='', sampling_types=['intra_landscape'], batch_sizes=[8, 64], criterion=None, device='cpu', verbose=False, random_data=False, number_classes=10) -> None: """ This method tests the scalar product between the teleporation line and the gradient, as well as between a random vector and the gradient for nullity. It then displays the histograms of the calculated scalar products. The method also aggregates all relevant micro teleportation data in a dataframe. Args: network : the model which we wish to use to compute the micro-teleporations dataset: the dataset that will be used to calculate the gradient and get dimensions for the neural teleportation model nb_teleport: The number of time the network is teleported and the scalar product calculated. An average is then calculated. network_descriptor: String describing the content of the network sampling_types : Teleportation sampling types, governs how the change of basis is computed batch_sizes: Size of the minibatch used to perform gradient calculation criterion: the loss function used to compute the gradient device: Device used to compute the network operations ('cpu' or 'cuda') verbose: If true, the method will output extensive details about the calculated vectors and aggregated data (mainly for debugging purposes) random_data: If True, random data with random labels is used for computing the gradient. If False, the dataset is used for computing the gradient. number_classes: Number of classes of the classification problem. """ # Arbitrary precision threshold for nullity comparison torch.set_printoptions(precision=10, sci_mode=True) tol = 1e-2 cobs = [0.001] hist_dir = f'images/histograms/{network_descriptor}' if torch.cuda.is_available(): print(f'{green}Using CUDA{reset}') network = network.cuda() if (criterion is None): loss_func = torch.nn.CrossEntropyLoss() else: loss_func = criterion # Initialize the dataframe for data aggregation aggregator = pd.DataFrame(columns=[ 'model name', 'sampling type', 'batch size', 'COB range', 'weights vector length', 'Micro-teleportation vs Gradient', 'Micro-teleportation vs Gradient std', 'Gradient vs Random Vector', 'Gradient vs Random Vector std', 'Random Vector vs Random Vector', 'Random Vector vs Random Vector std', 'Micro-teleportation vs Random Vector', 'Micro-teleportation vs Random Vector std' ]) for sampling_type in sampling_types: for batch_size in batch_sizes: dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size) data, target = next(iter(dataloader)) # save the initial weights for further reset model = NeuralTeleportationModel(network=network, input_shape=data.shape) if torch.cuda.is_available(): model = model.cuda() else: model = model.cpu() if torch.cuda.is_available(): w1 = model.get_weights().detach() else: w1 = model.get_weights().detach().numpy() for cob in cobs: angle_results = [] rand_angle_results = [] rand_rand_angle_results = [] rand_micro_angle_results = [] iterations = min( int(len(dataloader.dataset) / dataloader.batch_size), nb_teleport) for _ in tqdm(range(0, iterations)): # Get next data batch data, target = next(iter(dataloader)) if random_data: data, target = torch.rand(data.shape), torch.randint( 0, number_classes, target.shape) data, target = data.to(device), target.to(device) grad = model.get_grad(data, target, loss_func, zero_grad=False) # reset the weights for next teleportation model.set_weights(torch.tensor(w1)) # teleport and get the new weights model = model.random_teleport(cob_range=cob, sampling_type=sampling_type) if torch.cuda.is_available(): w2 = model.get_weights().detach() else: w2 = model.get_weights().detach().numpy() # get teleportation vector micro_teleport_vec = (w2 - w1) random_vector = torch.rand(grad.shape, dtype=torch.float) - 0.5 random_vector2 = torch.rand(grad.shape, dtype=torch.float) - 0.5 random_vector = random_vector.to(device) random_vector2 = random_vector2.to(device) # Normalized scalar products & angles calculations dot_prod = normalized_dot_product(grad, micro_teleport_vec) angle = np.degrees(torch.acos(dot_prod).cpu()) rand_dot_prod = normalized_dot_product(grad, random_vector) rand_angle = np.degrees(torch.acos(rand_dot_prod).cpu()) rand_rand_dot_prod = normalized_dot_product( random_vector2, random_vector) rand_rand_angle = np.degrees( torch.acos(rand_rand_dot_prod).cpu()) rand_micro_dot_prod = normalized_dot_product( random_vector2, micro_teleport_vec) rand_micro_angle = np.degrees( torch.acos(rand_micro_dot_prod).cpu()) # Perpendicularity assertion failed = (not torch.allclose( dot_prod, torch.tensor([0.0]).to(device), atol=tol)) rand_failed = (not torch.allclose(rand_dot_prod, torch.tensor( [0.0]).to(device), atol=tol)) target_angle = 90.0 angle_results.append(angle) rand_angle_results.append(rand_angle) rand_rand_angle_results.append(rand_rand_angle) rand_micro_angle_results.append(rand_micro_angle) angle_results = np.array(angle_results) rand_angle_results = np.array(rand_angle_results) rand_rand_angle_results = np.array(rand_rand_angle_results) rand_micro_angle_results = np.array(rand_micro_angle_results) # Append resuslts to dataframe for further ploting aggregator = aggregator.append( { 'model name': network_descriptor, 'sampling type': sampling_type, 'batch size': batch_size, 'COB range': cob, 'weights vector length': len(w1), 'Micro-teleportation vs Gradient': angle_results.mean(), 'Micro-teleportation vs Gradient std': angle_results.std(), 'Gradient vs Random Vector': rand_angle_results.mean(), 'Gradient vs Random Vector std': rand_angle_results.std(), 'Random Vector vs Random Vector': rand_rand_angle_results.mean(), 'Random Vector vs Random Vector std': rand_rand_angle_results.std(), 'Micro-teleportation vs Random Vector': rand_micro_angle_results.mean(), 'Micro-teleportation vs Random Vector std': rand_micro_angle_results.std() }, ignore_index=True) print( f'The angle between the gradient and a micro-teleporation vector is: ' f'{red * failed}' f'{np.round(angle_results.mean(), abs(int(np.log10(tol))))}', f' (!=0 => FAILED!)' * failed, f'{reset}', f' using {sampling_type} sampling type', f', the delta in angle is {angle - target_angle}°\n', f'The angle between the gradient and a random vector is: ', f'{red * rand_failed}{rand_angle_results.mean()}', f' (FAILED!)' * rand_failed, f'{reset}', f', the delta in angle is {rand_angle - target_angle}°\n', sep='') if verbose: print(aggregator.iloc[aggregator.last_valid_index()]) if torch.cuda.is_available(): print(f'w1: {w1}', f'nans: {torch.sum(torch.isnan(w1))}', f'max: {torch.max(w1)}', f'min: {torch.min(w1)}', sep='\n') print(f'w2: {w2}', f' nans: {torch.sum(torch.isnan(w2))}', f'max: {torch.max(w2)}', f'min: {torch.min(w2)}', sep='\n') else: print(f'w1: {w1}', f'nans: {np.sum(np.isnan(w1))}', f'max: {np.max(w1)}', f'min: {np.min(w1)}', sep='\n') print(f'w2: {w2}', f' nans: {np.sum(np.isnan(w2))}', f'max: {np.max(w2)}', f'min: {np.min(w2)}', sep='\n') if not np.isnan( aggregator.loc[aggregator.last_valid_index(), 'Micro-teleportation vs Gradient']): delta = 0.25 x_min = 90 - delta x_max = 90 + delta figsize = (10.0, 10.0) fig, (ax0, ax1, ax2, ax3) = plt.subplots(4, 1, figsize=figsize) if random_data: fig.suptitle( f'{network_descriptor} on Random Data and batch size of {batch_size}' ) else: fig.suptitle( f'{network_descriptor} on CIFAR-10 and batch size of {batch_size}' ) bin_height, bin_boundary = np.histogram( np.array(angle_results)) width = bin_boundary[1] - bin_boundary[0] bin_height = bin_height / float(max(bin_height)) ax0.bar(bin_boundary[:-1], bin_height, width=np.maximum(width, 0.01)) ax0.legend(['Micro-teleportation\n vs \n Gradient']) ax0.set_xlim(x_min, x_max) ax0.set_yticks([]) bin_height, bin_boundary = np.histogram( np.array(rand_micro_angle_results)) width = bin_boundary[1] - bin_boundary[0] bin_height = bin_height / float(max(bin_height)) ax1.bar(bin_boundary[:-1], bin_height, width=np.maximum(width, 0.1), color='g') ax1.set_xlim(x_min, x_max) ax1.legend(['Micro-teleportation\n vs \n Random Vector']) ax1.set_yticks([]) bin_height, bin_boundary = np.histogram( np.array(rand_angle_results)) width = bin_boundary[1] - bin_boundary[0] bin_height = bin_height / float(max(bin_height)) ax2.bar(bin_boundary[:-1], bin_height, width=np.maximum(width, 0.1), color='g') ax2.set_xlim(x_min, x_max) ax2.legend(['Gradient\n vs \n Random Vector']) ax2.set_yticks([]) bin_height, bin_boundary = np.histogram( np.array(rand_rand_angle_results)) width = bin_boundary[1] - bin_boundary[0] bin_height = bin_height / float(max(bin_height)) ax3.bar(bin_boundary[:-1], bin_height, width=np.maximum(width, 0.1), color='g') ax3.set_xlim(x_min, x_max) ax3.legend(['Random Vector\n vs \n Random Vector']) ax3.set_yticks([]) plt.xlabel('Angle in degrees') Path(hist_dir).mkdir(parents=True, exist_ok=True) plt.savefig( f'{hist_dir}/{network_descriptor}_' f'_cob_{cob}_iter_{iterations}_batch_size_{batch_size}.png' ) plt.show() if random_data: fig.savefig( f"{network_descriptor}-RandomData-batchsize_{batch_size}.pdf", bbox_inches='tight') else: fig.savefig( f"{network_descriptor}-cifar10-batchsize_{batch_size}.pdf", bbox_inches='tight') else: print(red) print(aggregator.iloc[aggregator.last_valid_index()]) print(reset)
def dot_product_between_teleportation(network, dataset, network_descriptor=None, nb_teleport=100, device='cpu') -> None: """ This method tests the scalar product between the initial and teleported set of weights and plots the results with respect to the order of magnitude of the change of basis of the teleportation Args: network : the model which we want to use to compute the teleportations dataset : the model which we want to use to size the teleportation model network_descriptor: String describing the content of the network nb_teleport: Number of times the micro-teleportation for statistical (mean, variance, etc) calculation device: Device used to compute the network operations ('cpu' or 'cuda') """ series_dir = f'images/series_dot_prod_vs_cob/{network_descriptor}' if torch.cuda.is_available(): print(f'{green}Using CUDA{reset}') network = network.cuda() if network_descriptor is None: network_descriptor = network.__name__ # Prepare the range of COB to test cobs = np.linspace(0.00001, 0.999, 40) dataloader = torch.utils.data.DataLoader(dataset, batch_size=16) data, target = next(iter(dataloader)) model = NeuralTeleportationModel(network=network, input_shape=data.shape) if torch.cuda.is_available(): model = model.cuda() else: model = model.cpu() w1 = model.get_weights().detach().to(device) dot_product_results = [] angles = [] for cob in cobs: dot_product_result = 0 angle = 0 for _ in tqdm(range(0, nb_teleport)): # reset the weights model.set_weights(w1) # teleport and get the new weights model.random_teleport(cob_range=cob, sampling_type='intra_landscape') w2 = model.get_weights().detach().to(device) # cos(theta) = (w1 w2)/(||w1|| ||w2||) dot_product_result += normalized_dot_product(w1, w2) angle += np.degrees( torch.acos(normalized_dot_product(w1, w2)).cpu()) dot_product_result /= nb_teleport angle /= nb_teleport dot_product_results.append(dot_product_result.item()) angles.append(angle.item()) plt.plot(cobs, dot_product_results) plt.title(f'Scalar product between original and \nteleported weights with ' f'respect to COB\'s order of magnitude\n{network_descriptor}') plt.ylabel('Scalar product') plt.xlabel('change of basis') Path(series_dir).mkdir(parents=True, exist_ok=True) plt.savefig( f'{series_dir}/dot_product_vs_cob_{network_descriptor}_Samp_type_intra_landscape' ) plt.show() plt.plot(cobs, angles) plt.title(f'Angle between original and \nteleported weights with ' f'respect to COB\'s order of magnitude\n{network_descriptor}') plt.ylabel('Theta') plt.xlabel('change of basis') Path(series_dir).mkdir(parents=True, exist_ok=True) plt.savefig( f'{series_dir}/angle_vs_cob_{network_descriptor}_Samp_type_intra_landscape' ) plt.show()