# MobileNetV2: Inverted Residuals and Linear Bottlenecks, CVPR 2018 from torch import nn from ..initialization import initialize_resnet from ..SharedUtils import additive_func, parse_channel_info class ConvBNReLU(nn.Module): def __init__(self, in_planes, out_planes, kernel_size, stride, groups, has_bn=True, has_relu=True): super(ConvBNReLU, self).__init__() padding = (kernel_size - 1) // 2 self.conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False) if has_bn: self.bn = nn.BatchNorm2d(out_planes) else : self.bn = None if has_relu: self.relu = nn.ReLU6(inplace=True) else : self.relu = None def forward(self, x): out = self.conv( x ) if self.bn: out = self.bn ( out ) if self.relu: out = self.relu( out ) return out class InvertedResidual(nn.Module): def __init__(self, channels, stride, expand_ratio, additive): super(InvertedResidual, self).__init__() self.stride = stride assert stride in [1, 2], 'invalid stride : {:}'.format(stride) assert len(channels) in [2, 3], 'invalid channels : {:}'.format(channels) if len(channels) == 2: layers = [] else: layers = [ConvBNReLU(channels[0], channels[1], 1, 1, 1)] layers.extend([ # dw ConvBNReLU(channels[-2], channels[-2], 3, stride, channels[-2]), # pw-linear ConvBNReLU(channels[-2], channels[-1], 1, 1, 1, True, False), ]) self.conv = nn.Sequential(*layers) self.additive = additive if self.additive and channels[0] != channels[-1]: self.shortcut = ConvBNReLU(channels[0], channels[-1], 1, 1, 1, True, False) else: self.shortcut = None self.out_dim = channels[-1] def forward(self, x): out = self.conv(x) # if self.additive: return additive_func(out, x) if self.shortcut: return out + self.shortcut(x) else : return out class InferMobileNetV2(nn.Module): def __init__(self, num_classes, xchannels, xblocks, dropout): super(InferMobileNetV2, self).__init__() block = InvertedResidual inverted_residual_setting = [ # t, c, n, s [1, 16 , 1, 1], [6, 24 , 2, 2], [6, 32 , 3, 2], [6, 64 , 4, 2], [6, 96 , 3, 1], [6, 160, 3, 2], [6, 320, 1, 1], ] assert len(inverted_residual_setting) == len(xblocks), 'invalid number of layers : {:} vs {:}'.format(len(inverted_residual_setting), len(xblocks)) for block_num, ir_setting in zip(xblocks, inverted_residual_setting): assert block_num <= ir_setting[2], '{:} vs {:}'.format(block_num, ir_setting) xchannels = parse_channel_info(xchannels) #for i, chs in enumerate(xchannels): # if i > 0: assert chs[0] == xchannels[i-1][-1], 'Layer[{:}] is invalid {:} vs {:}'.format(i, xchannels[i-1], chs) self.xchannels = xchannels self.message = 'InferMobileNetV2 : xblocks={:}'.format(xblocks) # building first layer features = [ConvBNReLU(xchannels[0][0], xchannels[0][1], 3, 2, 1)] last_channel_idx = 1 # building inverted residual blocks for stage, (t, c, n, s) in enumerate(inverted_residual_setting): for i in range(n): stride = s if i == 0 else 1 additv = True if i > 0 else False module = block(self.xchannels[last_channel_idx], stride, t, additv) features.append(module) self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, Cs={:}, stride={:}, expand={:}, original-C={:}".format(stage, i, n, len(features), self.xchannels[last_channel_idx], stride, t, c) last_channel_idx += 1 if i + 1 == xblocks[stage]: out_channel = module.out_dim for iiL in range(i+1, n): last_channel_idx += 1 self.xchannels[last_channel_idx][0] = module.out_dim break # building last several layers features.append(ConvBNReLU(self.xchannels[last_channel_idx][0], self.xchannels[last_channel_idx][1], 1, 1, 1)) assert last_channel_idx + 2 == len(self.xchannels), '{:} vs {:}'.format(last_channel_idx, len(self.xchannels)) # make it nn.Sequential self.features = nn.Sequential(*features) # building classifier self.classifier = nn.Sequential( nn.Dropout(dropout), nn.Linear(self.xchannels[last_channel_idx][1], num_classes), ) # weight initialization self.apply( initialize_resnet ) def get_message(self): return self.message def forward(self, inputs): features = self.features(inputs) vectors = features.mean([2, 3]) predicts = self.classifier(vectors) return features, predicts