| # 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) |
