Skip to content

Commit e01ee71

Browse files
hojelhojel
committed
add cfg converter
1 parent 2746474 commit e01ee71

File tree

3 files changed

+355
-19
lines changed

3 files changed

+355
-19
lines changed

‎create_yolo_caffemodel.py‎

Lines changed: 58 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,6 @@
99
import numpy as np
1010
import sys, getopt
1111

12-
def transpose_matrix(inputWeight, rows, cols):
13-
inputWeight_t = np.zeros((rows*cols,1))
14-
for x in xrange(rows):
15-
for y in xrange(cols):
16-
inputWeight_t[y*rows + x] = inputWeight[x*cols + y]
17-
return inputWeight_t
18-
1912
def main(argv):
2013
model_filename = ''
2114
yoloweight_filename = ''
@@ -53,20 +46,67 @@ def main(argv):
5346
print netWeights.shape
5447
count = 0
5548
for pr in params:
56-
biasSize = np.prod(net.params[pr][1].data.shape)
57-
net.params[pr][1].data[...] = np.reshape(netWeights[count:count+biasSize], net.params[pr][1].data.shape)
58-
count = count + biasSize
59-
weightSize = np.prod(net.params[pr][0].data.shape)
60-
if pr[0:2]=='co': # convolutional layer
61-
net.params[pr][0].data[...] = np.reshape(netWeights[count:count+weightSize], net.params[pr][0].data.shape)
62-
else: # fc layer
49+
lidx = list(net._layer_names).index(pr)
50+
layer = net.layers[lidx]
51+
if count == netWeights.shape[0] and (layer.type != 'BatchNorm' and layer.type != 'Scale'):
52+
print "WARNING: no weights left for %s" % pr
53+
break
54+
if layer.type == 'Convolution':
55+
print pr+"(conv)"
56+
# bias
57+
if len(net.params[pr]) > 1:
58+
bias_dim = net.params[pr][1].data.shape
59+
else:
60+
bias_dim = (net.params[pr][0].data.shape[0], )
61+
biasSize = np.prod(bias_dim)
62+
conv_bias = np.reshape(netWeights[count:count+biasSize], bias_dim)
63+
if len(net.params[pr]) > 1:
64+
assert(bias_dim == net.params[pr][1].data.shape)
65+
net.params[pr][1].data[...] = conv_bias
66+
conv_bias = None
67+
count = count + biasSize
68+
# batch_norm
69+
next_layer = net.layers[lidx+1]
70+
if next_layer.type == 'BatchNorm':
71+
bn_dims = (3, net.params[pr][0].data.shape[0])
72+
bnSize = np.prod(bn_dims)
73+
batch_norm = np.reshape(netWeights[count:count+bnSize], bn_dims)
74+
count = count + bnSize
75+
# weights
76+
dims = net.params[pr][0].data.shape
77+
weightSize = np.prod(dims)
78+
net.params[pr][0].data[...] = np.reshape(netWeights[count:count+weightSize], dims)
79+
count = count + weightSize
80+
elif layer.type == 'InnerProduct':
81+
print pr+"(fc)"
82+
# bias
83+
biasSize = np.prod(net.params[pr][1].data.shape)
84+
net.params[pr][1].data[...] = np.reshape(netWeights[count:count+biasSize], net.params[pr][1].data.shape)
85+
count = count + biasSize
86+
# weights
6387
dims = net.params[pr][0].data.shape
64-
if transFlag: # need transpose for fc layers
65-
net.params[pr][0].data[...] = np.reshape(transpose_matrix(netWeights[count:count+weightSize], dims[1],dims[0]), dims)
88+
weightSize = np.prod(dims)
89+
if transFlag:
90+
net.params[pr][0].data[...] = np.reshape(netWeights[count:count+weightSize], (dims[1], dims[0])).transpose()
6691
else:
6792
net.params[pr][0].data[...] = np.reshape(netWeights[count:count+weightSize], dims)
68-
count = count + weightSize
69-
print count
93+
count = count + weightSize
94+
elif layer.type == 'BatchNorm':
95+
print pr+"(batchnorm)"
96+
net.params[pr][0].data[...] = batch_norm[1] # mean
97+
net.params[pr][1].data[...] = batch_norm[2] # variance
98+
net.params[pr][2].data[...] = 1.0 # scale factor
99+
elif layer.type == 'Scale':
100+
print pr+"(scale)"
101+
net.params[pr][0].data[...] = batch_norm[0] # scale
102+
batch_norm = None
103+
if len(net.params[pr]) > 1:
104+
net.params[pr][1].data[...] = conv_bias # bias
105+
conv_bias = None
106+
else:
107+
print "WARNING: unsupported layer, "+pr
108+
if np.prod(netWeights.shape) != count:
109+
print "ERROR: size mismatch: %d" % count
70110
net.save(caffemodel_filename)
71111

72112
if __name__=='__main__':

‎create_yolo_prototxt.py‎

Lines changed: 296 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,296 @@
1+
# -*- coding: utf-8 -*-
2+
from ConfigParser import ConfigParser
3+
from collections import OrderedDict
4+
import logging
5+
import os
6+
import sys
7+
8+
class CaffeLayerGenerator(object):
9+
def __init__(self, name, ltype):
10+
self.name = name
11+
self.bottom = []
12+
self.top = []
13+
self.type = ltype
14+
def get_template(self):
15+
return """
16+
layer {{{{
17+
name: "{}"
18+
type: "{}"
19+
bottom: "{}"
20+
top: "{}"{{}}
21+
}}}}""".format(self.name, self.type, self.bottom[0], self.top[0])
22+
23+
class CaffeInputLayer(CaffeLayerGenerator):
24+
def __init__(self, name, channels, width, height):
25+
super(CaffeInputLayer, self).__init__(name, 'Input')
26+
self.channels = channels
27+
self.width = width
28+
self.height = height
29+
def write(self, f):
30+
f.write("""
31+
input: "{}"
32+
input_shape {{
33+
dim: 1
34+
dim: {}
35+
dim: {}
36+
dim: {}
37+
}}""".format(self.name, self.channels, self.width, self.height))
38+
39+
class CaffeConvolutionLayer(CaffeLayerGenerator):
40+
def __init__(self, name, filters, ksize=None, stride=None, pad=None, bias=True):
41+
super(CaffeConvolutionLayer, self).__init__(name, 'Convolution')
42+
self.filters = filters
43+
self.ksize = ksize
44+
self.stride = stride
45+
self.pad = pad
46+
self.bias = bias
47+
def write(self, f):
48+
opts = ['']
49+
if self.ksize is not None: opts.append('kernel_size: {}'.format(self.ksize))
50+
if self.stride is not None: opts.append('stride: {}'.format(self.stride))
51+
if self.pad is not None: opts.append('pad: {}'.format(self.pad))
52+
if not self.bias: opts.append('bias_term: false')
53+
param_str = """
54+
convolution_param {{
55+
num_output: {}{}
56+
}}""".format(self.filters, '\n '.join(opts))
57+
f.write(self.get_template().format(param_str))
58+
59+
class CaffePoolingLayer(CaffeLayerGenerator):
60+
def __init__(self, name, pooltype, ksize=None, stride=None, pad=None, global_pooling=None):
61+
super(CaffePoolingLayer, self).__init__(name, 'Pooling')
62+
self.pooltype = pooltype
63+
self.ksize = ksize
64+
self.stride = stride
65+
self.pad = pad
66+
self.global_pooling = global_pooling
67+
def write(self, f):
68+
opts = ['']
69+
if self.ksize is not None: opts.append('kernel_size: {}'.format(self.ksize))
70+
if self.stride is not None: opts.append('stride: {}'.format(self.stride))
71+
if self.pad is not None: opts.append('pad: {}'.format(self.pad))
72+
if self.global_pooling is not None: opts.append('global_pooling: {}'.format('True' if self.global_pooling else 'False'))
73+
param_str = """
74+
pooling_param {{
75+
pool: {}{}
76+
}}""".format(self.pooltype, '\n '.join(opts))
77+
f.write(self.get_template().format(param_str))
78+
79+
class CaffeInnerProductLayer(CaffeLayerGenerator):
80+
def __init__(self, name, num_output):
81+
super(CaffeInnerProductLayer, self).__init__(name, 'InnerProduct')
82+
self.num_output = num_output
83+
def write(self, f):
84+
param_str = """
85+
inner_product_param {{
86+
num_output: {}
87+
}}""".format(self.num_output)
88+
f.write(self.get_template().format(param_str))
89+
90+
class CaffeBatchNormLayer(CaffeLayerGenerator):
91+
def __init__(self, name):
92+
super(CaffeBatchNormLayer, self).__init__(name, 'BatchNorm')
93+
def write(self, f):
94+
param_str = """
95+
batch_norm_param {
96+
use_global_stats: true
97+
}"""
98+
f.write(self.get_template().format(param_str))
99+
100+
class CaffeScaleLayer(CaffeLayerGenerator):
101+
def __init__(self, name):
102+
super(CaffeScaleLayer, self).__init__(name, 'Scale')
103+
def write(self, f):
104+
param_str = """
105+
scale_param {
106+
bias_term: true
107+
}"""
108+
f.write(self.get_template().format(param_str))
109+
110+
class CaffeReluLayer(CaffeLayerGenerator):
111+
def __init__(self, name, negslope=None):
112+
super(CaffeReluLayer, self).__init__(name, 'Relu')
113+
self.negslope = negslope
114+
def write(self, f):
115+
param_str = ""
116+
if self.negslope is not None:
117+
param_str = """
118+
relu_param {{
119+
negative_slope: {}
120+
}}""".format(self.negslope)
121+
f.write(self.get_template().format(param_str))
122+
123+
class CaffeDropoutLayer(CaffeLayerGenerator):
124+
def __init__(self, name, prob):
125+
super(CaffeDropoutLayer, self).__init__(name, 'Dropout')
126+
self.prob = prob
127+
def write(self, f):
128+
param_str = """
129+
dropout_param {{
130+
dropout_ratio: {}
131+
}}""".format(self.prob)
132+
f.write(self.get_template().format(param_str))
133+
134+
class CaffeSoftmaxLayer(CaffeLayerGenerator):
135+
def __init__(self, name):
136+
super(CaffeSoftmaxLayer, self).__init__(name, 'Softmax')
137+
def write(self, f):
138+
f.write(self.get_template().format(""))
139+
140+
class CaffeProtoGenerator:
141+
def __init__(self, name):
142+
self.name = name
143+
self.sections = []
144+
self.lnum = 0
145+
self.layer = None
146+
def add_layer(self, l):
147+
self.sections.append( l )
148+
def add_input_layer(self, items):
149+
self.lnum = 0
150+
lname = "data"
151+
self.layer = CaffeInputLayer(lname, items['channels'], items['width'], items['height'])
152+
self.layer.top.append( lname )
153+
self.add_layer( self.layer )
154+
def add_convolution_layer(self, items):
155+
self.lnum += 1
156+
prev_blob = self.layer.top[0]
157+
lname = "conv"+str(self.lnum)
158+
filters = items['filters']
159+
ksize = items['size'] if 'size' in items else None
160+
stride = items['stride'] if 'stride' in items else None
161+
pad = items['pad'] if 'pad' in items else None
162+
bias = not bool(items['batch_normalize']) if 'batch_normalize' in items else True
163+
self.layer = CaffeConvolutionLayer( lname, filters, ksize=ksize, stride=stride, pad=pad, bias=bias )
164+
self.layer.bottom.append( prev_blob )
165+
self.layer.top.append( lname )
166+
self.add_layer( self.layer )
167+
def add_innerproduct_layer(self, items):
168+
self.lnum += 1
169+
prev_blob = self.layer.top[0]
170+
lname = "fc"+str(self.lnum)
171+
num_output = items['output']
172+
self.layer = CaffeInnerProductLayer( lname, num_output )
173+
self.layer.bottom.append( prev_blob )
174+
self.layer.top.append( lname )
175+
self.add_layer( self.layer )
176+
def add_pooling_layer(self, ltype, items, global_pooling=None):
177+
prev_blob = self.layer.top[0]
178+
lname = "pool"+str(self.lnum)
179+
ksize = items['size'] if 'size' in items else None
180+
stride = items['stride'] if 'stride' in items else None
181+
pad = items['pad'] if 'pad' in items else None
182+
self.layer = CaffePoolingLayer( lname, ltype, ksize=ksize, stride=stride, pad=pad, global_pooling=global_pooling )
183+
self.layer.bottom.append( prev_blob )
184+
self.layer.top.append( lname )
185+
self.add_layer( self.layer )
186+
def add_batchnorm_layer(self, items):
187+
prev_blob = self.layer.top[0]
188+
lname = "bn"+str(self.lnum)
189+
self.layer = CaffeBatchNormLayer( lname )
190+
self.layer.bottom.append( prev_blob )
191+
self.layer.top.append( lname )
192+
self.add_layer( self.layer )
193+
def add_scale_layer(self, items):
194+
prev_blob = self.layer.top[0]
195+
lname = "scale"+str(self.lnum)
196+
self.layer = CaffeScaleLayer( lname )
197+
self.layer.bottom.append( prev_blob )
198+
self.layer.top.append( lname )
199+
self.add_layer( self.layer )
200+
def add_relu_layer(self, items):
201+
prev_blob = self.layer.top[0]
202+
lname = "relu"+str(self.lnum)
203+
self.layer = CaffeReluLayer( lname )
204+
self.layer.bottom.append( prev_blob )
205+
self.layer.top.append( prev_blob ) # loopback
206+
self.add_layer( self.layer )
207+
def add_dropout_layer(self, items):
208+
prev_blob = self.layer.top[0]
209+
lname = "drop"+str(self.lnum)
210+
self.layer = CaffeDropoutLayer( lname, items['probability'] )
211+
self.layer.bottom.append( prev_blob )
212+
self.layer.top.append( prev_blob ) # loopback
213+
self.add_layer( self.layer )
214+
def add_softmax_layer(self, items):
215+
prev_blob = self.layer.top[0]
216+
lname = "prob"
217+
self.layer = CaffeSoftmaxLayer( lname )
218+
self.layer.bottom.append( prev_blob )
219+
self.layer.top.append( lname )
220+
self.add_layer( self.layer )
221+
def finalize(self, name):
222+
self.layer.top[0] = name # replace
223+
def write(self, fname):
224+
with open(fname, 'w') as f:
225+
f.write('name: "{}"'.format(self.name))
226+
for sec in self.sections:
227+
sec.write(f)
228+
logging.info('{} is generated'.format(fname))
229+
230+
###################################################################33
231+
class uniqdict(OrderedDict):
232+
_unique = 0
233+
def __setitem__(self, key, val):
234+
if isinstance(val, OrderedDict):
235+
self._unique += 1
236+
key += "_"+str(self._unique)
237+
OrderedDict.__setitem__(self, key, val)
238+
239+
def convert(cfgfile, ptxtfile):
240+
#
241+
parser = ConfigParser(dict_type=uniqdict)
242+
parser.read(cfgfile)
243+
netname = os.path.basename(cfgfile).split('.')[0]
244+
#print netname
245+
gen = CaffeProtoGenerator(netname)
246+
for section in parser.sections():
247+
_section = section.split('_')[0]
248+
if _section in ["crop", "cost"]:
249+
continue
250+
#
251+
batchnorm_followed = False
252+
relu_followed = False
253+
items = dict(parser.items(section))
254+
if 'batch_normalize' in items and items['batch_normalize']:
255+
batchnorm_followed = True
256+
if 'activation' in items and items['activation'] != 'linear':
257+
relu_followed = True
258+
#
259+
if _section == 'net':
260+
gen.add_input_layer(items)
261+
elif _section == 'convolutional':
262+
gen.add_convolution_layer(items)
263+
if batchnorm_followed:
264+
gen.add_batchnorm_layer(items)
265+
gen.add_scale_layer(items)
266+
if relu_followed:
267+
gen.add_relu_layer(items)
268+
elif _section == 'connected':
269+
gen.add_innerproduct_layer(items)
270+
if relu_followed:
271+
gen.add_relu_layer(items)
272+
elif _section == 'maxpool':
273+
gen.add_pooling_layer('MAX', items)
274+
elif _section == 'avgpool':
275+
gen.add_pooling_layer('AVE', items, global_pooling=True)
276+
elif _section == 'dropout':
277+
gen.add_dropout_layer(items)
278+
elif _section == 'softmax':
279+
gen.add_softmax_layer(items)
280+
else:
281+
logging.error("{} layer is not supported".format(_section))
282+
#gen.finalize('result')
283+
gen.write(ptxtfile)
284+
285+
def main():
286+
parser = argparse.ArgumentParser(description='Convert YOLO cfg to Caffe prototxt')
287+
parser.add_argument('cfg', type=str, help='YOLO cfg')
288+
parser.add_argument('prototxt', type=str, help='Caffe prototxt')
289+
args = parser.parse_args()
290+
291+
convert(args.cfg, args.prototxt)
292+
293+
if __name__ == "__main__":
294+
main()
295+
296+
# vim:sw=4:ts=4:et

‎yolo_main.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ def interpret_output(output, img_width, img_height):
4444
filter_mat_boxes = np.nonzero(filter_mat_probs)
4545
boxes_filtered = boxes[filter_mat_boxes[0],filter_mat_boxes[1],filter_mat_boxes[2]]
4646
probs_filtered = probs[filter_mat_probs]
47-
classes_num_filtered = np.argmax(filter_mat_probs,axis=3)[filter_mat_boxes[0],filter_mat_boxes[1],filter_mat_boxes[2]]
47+
classes_num_filtered = np.argmax(probs,axis=3)[filter_mat_boxes[0],filter_mat_boxes[1],filter_mat_boxes[2]]
4848

4949
argsort = np.array(np.argsort(probs_filtered))[::-1]
5050
boxes_filtered = boxes_filtered[argsort]

0 commit comments

Comments
 (0)