# Module caffe2.python.models.shufflenet
from caffe2.python import brew
"""
Utilitiy for creating ShuffleNet
"ShuffleNet V2: Practical Guidelines for EfficientCNN Architecture Design" by Ma et. al. 2018
"""
OUTPUT_CHANNELS = {
'0.5x': [24, 48, 96, 192, 1024],
'1.0x': [24, 116, 232, 464, 1024],
'1.5x': [24, 176, 352, 704, 1024],
'2.0x': [24, 244, 488, 976, 2048],
}
class ShuffleNetV2Builder():
def __init__(
self,
model,
data,
num_input_channels,
num_labels,
num_groups=2,
width='1.0x',
is_test=False,
detection=False,
bn_epsilon=1e-5,
):
self.model = model
self.prev_blob = data
self.num_input_channels = num_input_channels
self.num_labels = num_labels
self.num_groups = num_groups
self.output_channels = OUTPUT_CHANNELS[width]
self.stage_repeats = [3, 7, 3]
self.is_test = is_test
self.detection = detection
self.bn_epsilon = bn_epsilon
def create(self):
in_channels = self.output_channels[0]
self.prev_blob = brew.conv(self.model, self.prev_blob, 'stage1_conv',
self.num_input_channels, in_channels,
weight_init=("MSRAFill", {}),
kernel=3, stride=2)
self.prev_blob = brew.max_pool(self.model, self.prev_blob,
'stage1_pool', kernel=3, stride=2)
# adds stage#{2,3,4}; see table 5 of the ShufflenetV2 paper.
for idx, (out_channels, n_repeats) in enumerate(zip(
self.output_channels[1:4], self.stage_repeats
)):
prefix = 'stage{}_stride{}'.format(idx + 2, 2)
self.add_spatial_ds_unit(prefix, in_channels, out_channels)
in_channels = out_channels
for i in range(n_repeats):
prefix = 'stage{}_stride{}_repeat{}'.format(
idx + 2, 1, i + 1
)
self.add_basic_unit(prefix, in_channels)
self.last_conv = brew.conv(self.model, self.prev_blob, 'conv5',
in_channels, self.output_channels[4],
kernel=1)
self.avg_pool = self.model.AveragePool(self.last_conv, 'avg_pool',
kernel=7)
self.last_out = brew.fc(self.model,
self.avg_pool,
'last_out_L{}'.format(self.num_labels),
self.output_channels[4],
self.num_labels)
# spatial down sampling unit with stride=2
def add_spatial_ds_unit(self, prefix, in_channels, out_channels, stride=2):
right = left = self.prev_blob
out_channels = out_channels // 2
# Enlarge the receptive field for detection task
if self.detection:
left = self.add_detection_unit(left, prefix + '_left_detection',
in_channels, in_channels)
left = self.add_dwconv3x3_bn(left, prefix + 'left_dwconv',
in_channels, stride)
left = self.add_conv1x1_bn(left, prefix + '_left_conv1', in_channels,
out_channels)
if self.detection:
right = self.add_detection_unit(right, prefix + '_right_detection',
in_channels, in_channels)
right = self.add_conv1x1_bn(right, prefix + '_right_conv1',
in_channels, out_channels)
right = self.add_dwconv3x3_bn(right, prefix + '_right_dwconv',
out_channels, stride)
right = self.add_conv1x1_bn(right, prefix + '_right_conv2',
out_channels, out_channels)
self.prev_blob = brew.concat(self.model, [right, left],
prefix + '_concat')
self.prev_blob = self.model.net.ChannelShuffle(
self.prev_blob, prefix + '_ch_shuffle',
group=self.num_groups, kernel=1
)
# basic unit with stride=1
def add_basic_unit(self, prefix, in_channels, stride=1):
in_channels = in_channels // 2
left = prefix + '_left'
right = prefix + '_right'
self.model.net.Split(self.prev_blob, [left, right])
if self.detection:
right = self.add_detection_unit(right, prefix + '_right_detection',
in_channels, in_channels)
right = self.add_conv1x1_bn(right, prefix + '_right_conv1',
in_channels, in_channels)
right = self.add_dwconv3x3_bn(right, prefix + '_right_dwconv',
in_channels, stride)
right = self.add_conv1x1_bn(right, prefix + '_right_conv2',
in_channels, in_channels)
self.prev_blob = brew.concat(self.model, [right, left],
prefix + '_concat')
self.prev_blob = self.model.net.ChannelShuffle(
self.prev_blob, prefix + '_ch_shuffle',
group=self.num_groups, kernel=1
)
# helper functions to create net's units
def add_detection_unit(self, prev_blob, prefix, in_channels, out_channels,
kernel=3, pad=1):
out_blob = brew.conv(self.model, prev_blob, prefix + '_conv',
in_channels, out_channels, kernel=kernel,
weight_init=("MSRAFill", {}),
group=in_channels, pad=pad)
out_blob = brew.spatial_bn(self.model, out_blob, prefix + '_bn',
out_channels, epsilon=self.bn_epsilon,
is_test=self.is_test)
return out_blob
def add_conv1x1_bn(self, prev_blob, blob, in_channels, out_channels):
prev_blob = brew.conv(self.model, prev_blob, blob, in_channels,
out_channels, kernel=1,
weight_init=("MSRAFill", {}))
prev_blob = brew.spatial_bn(self.model, prev_blob, prev_blob + '_bn',
out_channels,
epsilon=self.bn_epsilon,
is_test=self.is_test)
prev_blob = brew.relu(self.model, prev_blob, prev_blob)
return prev_blob
def add_dwconv3x3_bn(self, prev_blob, blob, channels, stride):
prev_blob = brew.conv(self.model, prev_blob, blob, channels,
channels, kernel=3,
weight_init=("MSRAFill", {}),
stride=stride, group=channels, pad=1)
prev_blob = brew.spatial_bn(self.model, prev_blob,
prev_blob + '_bn',
channels,
epsilon=self.bn_epsilon,
is_test=self.is_test)
return prev_blob
def create_shufflenet(
model,
data,
num_input_channels,
num_labels,
label=None,
is_test=False,
no_loss=False,
):
builder = ShuffleNetV2Builder(model, data, num_input_channels,
num_labels,
is_test=is_test)
builder.create()
if no_loss:
return builder.last_out
if (label is not None):
(softmax, loss) = model.SoftmaxWithLoss(
[builder.last_out, label],
["softmax", "loss"],
)
return (softmax, loss)