(self, block_args, global_params, image_size=None)
| 56 | """ |
| 57 | |
| 58 | def __init__(self, block_args, global_params, image_size=None): |
| 59 | super().__init__() |
| 60 | self._block_args = block_args |
| 61 | self._bn_mom = 1 - global_params.batch_norm_momentum # pytorch's difference from tensorflow |
| 62 | self._bn_eps = global_params.batch_norm_epsilon |
| 63 | self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1) |
| 64 | self.id_skip = block_args.id_skip # whether to use skip connection and drop connect |
| 65 | |
| 66 | # Expansion phase (Inverted Bottleneck) |
| 67 | inp = self._block_args.input_filters # number of input channels |
| 68 | oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels |
| 69 | if self._block_args.expand_ratio != 1: |
| 70 | Conv2d = get_same_padding_conv2d(image_size=image_size) |
| 71 | self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False) |
| 72 | self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) |
| 73 | # image_size = calculate_output_image_size(image_size, 1) <-- this wouldn't modify image_size |
| 74 | |
| 75 | # Depthwise convolution phase |
| 76 | k = self._block_args.kernel_size |
| 77 | s = self._block_args.stride |
| 78 | Conv2d = get_same_padding_conv2d(image_size=image_size) |
| 79 | self._depthwise_conv = Conv2d( |
| 80 | in_channels=oup, out_channels=oup, groups=oup, # groups makes it depthwise |
| 81 | kernel_size=k, stride=s, bias=False) |
| 82 | self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) |
| 83 | image_size = calculate_output_image_size(image_size, s) |
| 84 | |
| 85 | # Squeeze and Excitation layer, if desired |
| 86 | if self.has_se: |
| 87 | Conv2d = get_same_padding_conv2d(image_size=(1, 1)) |
| 88 | num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio)) |
| 89 | self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1) |
| 90 | self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1) |
| 91 | |
| 92 | # Pointwise convolution phase |
| 93 | final_oup = self._block_args.output_filters |
| 94 | Conv2d = get_same_padding_conv2d(image_size=image_size) |
| 95 | self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False) |
| 96 | self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps) |
| 97 | self._swish = MemoryEfficientSwish() |
| 98 | |
| 99 | def forward(self, inputs, drop_connect_rate=None): |
| 100 | """MBConvBlock's forward function. |
no test coverage detected