Mobile Inverted Residual Bottleneck Block Args: block_args (namedtuple): BlockArgs, see above global_params (namedtuple): GlobalParam, see above Attributes: has_se (bool): Whether the block contains a Squeeze and Excitation layer.
| 15 | ) |
| 16 | |
| 17 | class MBConvBlock(nn.Module): |
| 18 | """ |
| 19 | Mobile Inverted Residual Bottleneck Block |
| 20 | |
| 21 | Args: |
| 22 | block_args (namedtuple): BlockArgs, see above |
| 23 | global_params (namedtuple): GlobalParam, see above |
| 24 | |
| 25 | Attributes: |
| 26 | has_se (bool): Whether the block contains a Squeeze and Excitation layer. |
| 27 | """ |
| 28 | |
| 29 | def __init__(self, block_args, global_params): |
| 30 | super().__init__() |
| 31 | self._block_args = block_args |
| 32 | self._bn_mom = 1 - global_params.batch_norm_momentum |
| 33 | self._bn_eps = global_params.batch_norm_epsilon |
| 34 | self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1) |
| 35 | self.id_skip = block_args.id_skip # skip connection and drop connect |
| 36 | |
| 37 | # Get static or dynamic convolution depending on image size |
| 38 | Conv2d = get_same_padding_conv2d(image_size=global_params.image_size) |
| 39 | |
| 40 | # Expansion phase |
| 41 | inp = self._block_args.input_filters # number of input channels |
| 42 | oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels |
| 43 | if self._block_args.expand_ratio != 1: |
| 44 | self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False) |
| 45 | self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) |
| 46 | |
| 47 | # Depthwise convolution phase |
| 48 | k = self._block_args.kernel_size |
| 49 | s = self._block_args.stride |
| 50 | self._depthwise_conv = Conv2d( |
| 51 | in_channels=oup, out_channels=oup, groups=oup, # groups makes it depthwise |
| 52 | kernel_size=k, stride=s, bias=False) |
| 53 | self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) |
| 54 | |
| 55 | # Squeeze and Excitation layer, if desired |
| 56 | if self.has_se: |
| 57 | num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio)) |
| 58 | self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1) |
| 59 | self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1) |
| 60 | |
| 61 | # Output phase |
| 62 | final_oup = self._block_args.output_filters |
| 63 | self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False) |
| 64 | self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps) |
| 65 | self._swish = MemoryEfficientSwish() |
| 66 | |
| 67 | def forward(self, inputs, drop_connect_rate=None): |
| 68 | """ |
| 69 | :param inputs: input tensor |
| 70 | :param drop_connect_rate: drop connect rate (float, between 0 and 1) |
| 71 | :return: output of block |
| 72 | """ |
| 73 | |
| 74 | # Expansion and Depthwise Convolution |