Repository navigation
Expand file tree
/
Copy pathnet2net.py
More file actions
executable file
·262 lines (214 loc) · 9.61 KB
/
Copy pathnet2net.py
File metadata and controls
executable file
·262 lines (214 loc) · 9.61 KB
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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
#!/usr/bin/env python3
#
# This file is part of Leela Zero.
# Copyright (C) 2017 Henrik Forsten
#
# Leela Zero is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Leela Zero is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Leela Zero. If not, see <http://www.gnu.org/licenses/>.
import argparse
import os
import numpy as np
import scipy.signal as signal
from copy import deepcopy
def convolve(w, x , bn=None, bn_epsilon=1e-5):
x_ch, x_w, x_h = x.shape
outputs, inputs, _, __ = w.shape
assert x_ch == inputs
res = np.zeros((outputs, x_w, x_h))
for o in range(outputs):
for c in range(inputs):
res[o,:,:] += signal.correlate2d(x[c,:,:], w[o,c,:,:], mode='same')
if bn == None:
return res
bn_means = bn[0]
bn_vars = bn[1]
for o in range(outputs):
scale = 1.0 / np.sqrt(bn_epsilon + bn_vars[o])
v = scale * (res[o,:,:] - bn_means[o])
res[o,:,:] = np.maximum(v, 0)
return res
def read_net(filename):
with open(filename, 'r') as f:
weights = []
for e, line in enumerate(f):
if e == 0:
print("Version", line.strip())
if line != '1\n':
raise ValueError("Unknown version {}".format(line.strip()))
else:
weights.append(list(map(float, line.split(' '))))
if e == 2:
channels = len(line.split(' '))
print("Channels", channels)
blocks = e - (4 + 14)
if blocks % 8 != 0:
raise ValueError("Inconsistent number of weights in the file")
blocks //= 8
print("Blocks", blocks)
return blocks, channels, weights
def conv_bn_wider(weights, next_weights, inputs, channels,
new_channels, noise_std=0, last_block=False,
rand=None, dir_alpha=None, verify=False):
if new_channels == 0:
return weights, next_weights
if rand == None:
rand = list(range(channels))
rand.extend(np.random.randint(0, channels, new_channels))
rep_factor = np.bincount(rand)
factor = np.zeros(len(rand))
#In the net2net paper every input weight was weighted equally,
#but in general we can have unequal division of the weights
if dir_alpha == None:
#Equal division
for i in range(len(rand)):
factor[i] = 1.0/rep_factor[rand[i]]
else:
#Unequal input weighting determined by dirichlet distribution
for i in range(channels):
x = np.random.dirichlet([dir_alpha]*rep_factor[i])
e = 0
for j in range(channels + new_channels):
if rand[j] == i:
factor[j] = x[e]
e += 1
#Widen the current layer
w_conv_new = np.array(weights[0]).reshape(channels, inputs, 3, 3)[rand, :, :, :]
bias = np.array(weights[1])[rand]
w_bn_means = np.array(weights[2])[rand]
w_bn_vars = np.array(weights[3])[rand]
#Widen the next layer inputs
if not last_block:
w_filter = 3
else:
w_filter = 1
next_weights_new = []
for j in range(len(next_weights)):
n = np.array(next_weights[j]).reshape(-1, channels, w_filter, w_filter)
next_weights_new.append(n[:, rand, :, :])
for i in range(len(rand)):
noise = 0
if i >= channels:
noise = np.random.normal(0, noise_std)
next_weights_new[j][:, i, :, :] *= (1.0 + noise)*factor[i]
if noise_std == 0 and verify:
x = np.random.random((inputs, 19, 19))
old1 = convolve(np.array(weights[0]).reshape(channels, inputs, 3, 3), x, bn=[weights[2], weights[3]])
old2 = convolve(np.array(next_weights[0]).reshape(-1, channels, w_filter, w_filter), old1)
new1 = convolve(np.array(w_conv_new).reshape(channels + new_channels, inputs, 3, 3), x, bn=[w_bn_means, w_bn_vars])
new2 = convolve(np.array(next_weights_new[0]).reshape(-1, channels + new_channels, w_filter, w_filter), new1)
assert (np.abs(old2 - new2) < 1e-6).all()
w_conv_new = w_conv_new.flatten()
for j in range(len(next_weights)):
next_weights_new[j] = next_weights_new[j].flatten()
w_new = [w_conv_new, bias, w_bn_means, w_bn_vars]
return w_new, next_weights_new
def write_layer(weights, out_file):
for w in weights:
out_file.write(' '.join(map(str,w)) + '\n')
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Add filters/blocks to existing network such that the output is preserved.')
parser.add_argument("blocks", help="Residual blocks to add", type=int)
parser.add_argument("filters", help="Filters to add", type=int)
parser.add_argument("network", help="Input network", type=str)
parser.add_argument("--noise", nargs='?', help="Standard deviation of noise to add to new filters/blocks. Default: 5e-3",
default=5e-3, type=float)
parser.add_argument("--dir_alpha", nargs='?', help=\
"""Dirichlet distribution parameter for input weight distribution for replicated channels. """\
"""Larger values divide input values more equally. """\
"""Smaller ones give one large input weight while others are very small. """\
"""You probably want this to be at least 1 to avoid near zero weights. """\
"""Set to 0 to divide input weights equally. Default: 10""",
default=10, type=float)
parser.add_argument("--verify", help="Verify that output matches. Noise must be disabled.",
default=False, action='store_true')
parser.add_argument("--add_inputs", help="Adds input planes to network",
default=0, type=int)
args = parser.parse_args()
new_blocks = args.blocks
new_channels = args.filters
net_filename = args.network
noise_std = args.noise
dir_alpha = args.dir_alpha
verify = args.verify
if dir_alpha <= 0:
dir_alpha = None
if verify and noise_std != 0:
raise ValueError("Noise must be zero if verify is enabled.")
base, ext = os.path.splitext(net_filename)
output_filename = base + "_net2net" + ext
blocks, channels, weights = read_net(net_filename)
if new_blocks < 0:
raise ValueError("Blocks must be non-negative")
if new_channels < 0:
raise ValueError("Filters must be non-negative")
print("Output will have {} blocks and {} channels.".format(
blocks+new_blocks, channels+new_channels))
input_planes = 18
#Input convolution, bias, batch norm means, batch norm variances
w_input = weights[:4]
#Residual block convolution + batch norm
w_convs = []
for b in range(2*blocks):
w_convs.append(weights[4 + b*4: 4 + (b+1)*4])
i = ((b+1)*4) + 4
w_pol = weights[i:i+6]
w_val = weights[i+6:]
if new_blocks > 0:
#New blocks must have zero output due to the residual connection
new_block_conv = np.random.normal(0, noise_std, 9*(channels)**2)
new_block_bias = np.zeros(channels)
new_block_bn_mean = new_block_bias.copy()
new_block_bn_variances = np.ones(channels)
new_block = [new_block_conv, new_block_bias, new_block_bn_mean, new_block_bn_variances]
for i in range(2*new_blocks):
w_convs.append(deepcopy(new_block))
blocks += new_blocks
out_file = open(output_filename, 'w')
#Version
out_file.write('1\n')
#Making widening choice deterministic allows residual connection to be left
#as identity map. If the choice is not deterministic then the output of the
#widened network doesn't match the original one.
rand = list(range(channels))
rand.extend(np.random.randint(0, channels, new_channels))
if args.add_inputs > 0:
w_in_new = np.array(w_input[0]).reshape(channels, input_planes, 3, 3)
noise = np.random.normal(0, noise_std, [channels, args.add_inputs, 3, 3])
w_in_new = np.append(w_in_new, noise, axis=1)
input_planes = input_planes + args.add_inputs
print("Output will have {} input planes".format(input_planes))
w_input[0] = w_in_new.flatten()
#Input
w_wider, conv_next = conv_bn_wider(w_input, [w_convs[0][0]], input_planes,
channels, new_channels, noise_std, rand=rand, dir_alpha=dir_alpha, verify=verify)
w_convs[0][0] = conv_next[0]
write_layer(w_wider, out_file)
for e, w in enumerate(w_convs[:-1]):
r = rand
if e % 2 == 0:
print("Processing block", 1 + e//2)
#First convolution in residual block can be widened randomly
r = None
w_wider, conv_next = conv_bn_wider(w, [w_convs[e+1][0]], channels + new_channels,
channels, new_channels, noise_std, rand=r, dir_alpha=dir_alpha, verify=verify)
w_convs[e+1][0] = conv_next[0]
write_layer(w_wider, out_file)
#The last block is special case because of policy and value heads
w_wider, w_next = conv_bn_wider(w_convs[-1], [w_pol[0], w_val[0]], channels + new_channels,
channels, new_channels, noise_std, last_block=True, rand=rand, dir_alpha=dir_alpha, verify=verify)
w_pol[0] = w_next[0]
w_val[0] = w_next[1]
write_layer(w_wider, out_file)
write_layer(w_pol, out_file)
write_layer(w_val, out_file)
out_file.close()