-
Notifications
You must be signed in to change notification settings - Fork 3
/
calculate_cob_test.py
160 lines (119 loc) · 5.82 KB
/
calculate_cob_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import numpy as np
import torch
from neuralteleportation.neuralteleportationmodel import NeuralTeleportationModel
def test_calculate_cob(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 random teleportation.
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)
w1 = model.get_weights(concat=False, flatten=False, bias=False)
model.random_teleport()
w2 = model.get_weights(concat=False, flatten=False, bias=False)
if noise:
for w in w2:
w += torch.rand(w.shape) * 0.001
cob = model.get_cob()
calculated_cob = model.calculate_cob(w1, w2)
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.detach().numpy(), calculated_cob.detach().numpy(), atol=1e-6),\
"Calculate cob FAILED for " + model_name + " model."
print("Calculate cob successful for " + model_name + " model.")
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.")
if __name__ == '__main__':
import torch.nn as nn
from torch.nn.modules import Flatten
from neuralteleportation.layers.layer_utils import swap_model_modules_for_COB_modules
mlp_model = torch.nn.Sequential(
Flatten(),
nn.Linear(784, 128),
nn.ReLU(),
nn.Linear(128, 128),
nn.ReLU(),
nn.Linear(128, 128),
nn.ReLU(),
nn.Linear(128, 10)
)
mlp_model = swap_model_modules_for_COB_modules(mlp_model)
test_calculate_cob(network=mlp_model, model_name="MLP", verbose=True)
test_calculate_cob_weights(network=mlp_model, model_name="MLP", verbose=True)
test_calculate_ones(network=mlp_model, model_name="MLP", verbose=True)