| 21 | |
| 22 | |
| 23 | class ShuffleNetV2Builder(): |
| 24 | def __init__( |
| 25 | self, |
| 26 | model, |
| 27 | data, |
| 28 | num_input_channels, |
| 29 | num_labels, |
| 30 | num_groups=2, |
| 31 | width='1.0x', |
| 32 | is_test=False, |
| 33 | detection=False, |
| 34 | bn_epsilon=1e-5, |
| 35 | ): |
| 36 | self.model = model |
| 37 | self.prev_blob = data |
| 38 | self.num_input_channels = num_input_channels |
| 39 | self.num_labels = num_labels |
| 40 | self.num_groups = num_groups |
| 41 | self.output_channels = OUTPUT_CHANNELS[width] |
| 42 | self.stage_repeats = [3, 7, 3] |
| 43 | self.is_test = is_test |
| 44 | self.detection = detection |
| 45 | self.bn_epsilon = bn_epsilon |
| 46 | |
| 47 | def create(self): |
| 48 | in_channels = self.output_channels[0] |
| 49 | |
| 50 | self.prev_blob = brew.conv(self.model, self.prev_blob, 'stage1_conv', |
| 51 | self.num_input_channels, in_channels, |
| 52 | weight_init=("MSRAFill", {}), |
| 53 | kernel=3, stride=2) |
| 54 | self.prev_blob = brew.max_pool(self.model, self.prev_blob, |
| 55 | 'stage1_pool', kernel=3, stride=2) |
| 56 | |
| 57 | # adds stage#{2,3,4}; see table 5 of the ShufflenetV2 paper. |
| 58 | for idx, (out_channels, n_repeats) in enumerate(zip( |
| 59 | self.output_channels[1:4], self.stage_repeats |
| 60 | )): |
| 61 | prefix = 'stage{}_stride{}'.format(idx + 2, 2) |
| 62 | self.add_spatial_ds_unit(prefix, in_channels, out_channels) |
| 63 | in_channels = out_channels |
| 64 | for i in range(n_repeats): |
| 65 | prefix = 'stage{}_stride{}_repeat{}'.format( |
| 66 | idx + 2, 1, i + 1 |
| 67 | ) |
| 68 | self.add_basic_unit(prefix, in_channels) |
| 69 | |
| 70 | self.last_conv = brew.conv(self.model, self.prev_blob, 'conv5', |
| 71 | in_channels, self.output_channels[4], |
| 72 | kernel=1) |
| 73 | self.avg_pool = self.model.AveragePool(self.last_conv, 'avg_pool', |
| 74 | kernel=7) |
| 75 | self.last_out = brew.fc(self.model, |
| 76 | self.avg_pool, |
| 77 | 'last_out_L{}'.format(self.num_labels), |
| 78 | self.output_channels[4], |
| 79 | self.num_labels) |
| 80 |
no outgoing calls
no test coverage detected
searching dependent graphs…