def ShuffleNetBody(net, from_layer, pool_ave=False, fc1000=False):
# parameters for convolution layer with batchnorm.
kwargs = {
'weight_filler': dict(type='msra'),
'bias_term': False,
}
# parameters for batchnorm layer.
bn_kwargs = {
'param': [
dict(lr_mult=0, decay_mult=0),
dict(lr_mult=0, decay_mult=0),
dict(lr_mult=0, decay_mult=0)],
}
sb_kwargs = {
'bias_term': True,
'filler': dict(value=1.0),
'bias_filler': dict(value=0),
}
assert from_layer in net.keys()
net.conv1 = L.Convolution(net[from_layer], num_output=24, pad=1, stride=2, kernel_size=3, **kwargs)
net.conv1_bn = L.BatchNorm(net.conv1, in_place=True, **bn_kwargs)
net.conv1_scale = L.Scale(net.conv1, in_place=True, **sb_kwargs)
net.conv1_relu = L.ReLU(net.conv1, in_place=True)
net.pool1 = L.Pooling(net.conv1, pool=P.Pooling.MAX, kernel_size=3, stride=2)
from_layer = "pool1"
strides=[[0],[2],[1],[1],[1],[2],[1],[1],[1],[1],[1],[1],[1],[2],[1],[1],[1]]
num_outputs=[24,58,58,58,58,116,116,116,116,116,116,116,116,232,232,232,232]
for i in range(1,17):
if i == 1 or i ==5 or i == 13:
from_layer1 = from_layer
for j in range(1,3):
out_layer = "branch{}_1_conv{}".format(i,j)
if j == 1:
if i == 1:
net[out_layer] = L.ConvolutionDepthwise(net[from_layer1], convolution_param=dict(kernel_size=3, stride=strides[i], num_output=num_outputs[i-1], pad=1, **kwargs))
else:
net[out_layer] = L.ConvolutionDepthwise(net[from_layer1], convolution_param=dict(kernel_size=3, stride=strides[i], num_output=num_outputs[i], pad=1, **kwargs))
else:
net[out_layer] = L.Convolution(net[from_layer1], num_output=num_outputs[i], kernel_size=1, stride=1, pad=0, **kwargs)
net[out_layer+'_bn'] = L.BatchNorm(net[out_layer], in_place=True, **bn_kwargs)
net[out_layer+'_scale'] = L.Scale(net[out_layer], in_place=True, **sb_kwargs)
if j == 2:
net[out_layer+'_relu'] = L.ReLU(net[out_layer], in_place=True)
from_layer1 = out_layer
else:
out_layer = 'slice{}'.format(i)
net[out_layer] = L.Slice(net[from_layer])
from_layer = out_layer
from_layer1 = from_layer
from_layer2 = from_layer
for j in range(1,4):
out_layer = "branch{}_2_conv{}".format(i,j)
if j == 2:
net[out_layer] = L.ConvolutionDepthwise(net[from_layer2], convolution_param=dict(kernel_size=3, stride=strides[i], num_output=num_outputs[i], pad=1, **kwargs))
else:
net[out_layer] = L.Convolution(net[from_layer2], num_output=num_outputs[i], kernel_size=1, stride=1, pad=0, **kwargs)
net[out_layer+'_bn'] = L.BatchNorm(net[out_layer], in_place=True, **bn_kwargs)
net[out_layer+'_scale'] = L.Scale(net[out_layer], in_place=True, **sb_kwargs)
if j == 1 or j == 3:
net[out_layer+'_relu'] = L.ReLU(net[out_layer], in_place=True)
from_layer2 = out_layer
from_layer = 'concat{}'.format(i)
out_layer = 'shuffle{}'.format(i)
net[from_layer] = L.Concat(net[from_layer1],net[from_layer2])
net[out_layer] = L.ShuffleChannel(net[from_layer],shuffle_channel_param=dict(group=2))
from_layer = out_layer
# add conv5
net.conv5 = L.Convolution(net[from_layer],num_output=1024,pad=0,stride=1, kernel_size=1, **kwargs)
net.conv5_bn = L.BatchNorm(net.conv5, in_place=True, **bn_kwargs)
net.conv5_scale = L.Scale(net.conv5, in_place=True, **sb_kwargs)
net.conv5_relu = L.ReLU(net.conv5, in_place=True)
from_layer = 'conv5'
if pool_ave:
net.pool_ave = L.Pooling(net[from_layer], pool=P.Pooling.AVE, global_pooling=True)
from_layer = pool_ave
if fc1000:
kwargs = {
'param': [dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)],
'weight_filler': dict(type='msra'),
'bias_filler': dict(type='constant', value=0)}
net.fc1000 = L.Convolution(net[from_layer], kernel_size=1,num_output=1000,**kargs)
return net