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_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.")