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 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.")
print("Initial: w1 - w2 ([:100]): ",
(model1.get_weights() - model2.get_weights()).abs()[:100])
print("Initial: w1 - w2 ([-100:]): ",
(model1.get_weights() - model2.get_weights()).abs()[-100:])
w1 = model1.get_weights()
w2 = model2.get_weights()
diff = (w1.detach().cpu() - w2.detach().cpu()).abs().mean()
print("Initial weight difference :", diff)
w1 = model1.get_weights(concat=False, flatten=False, bias=False)
w2 = model2.get_weights(concat=False, flatten=False, bias=False)
calculated_cob = model1.calculate_cob(w1,
w2,
concat=True,
eta=0.00001,
steps=6000)
torch.save(calculated_cob, pjoin(save_path, 'calculated_cob.pt'))
model1.teleport(calculated_cob)
w1 = model1.get_weights()
w2 = model2.get_weights()
diff = (w1.detach().cpu() - w2.detach().cpu()).abs().mean()
print("Predicted weight difference :", diff)
w1 = model1.get_weights(concat=False, flatten=False, bias=False)
w2 = model2.get_weights(concat=False, flatten=False, bias=False)
print("Weight difference by layer:")