Multi-scale convolution block (MSCB)
| 111 | return outputs |
| 112 | |
| 113 | class MSCB(nn.Module): |
| 114 | """ |
| 115 | Multi-scale convolution block (MSCB) |
| 116 | """ |
| 117 | def __init__(self, in_channels, out_channels, stride, kernel_sizes=[1,3,5], expansion_factor=2, dw_parallel=True, add=True, activation='relu6'): |
| 118 | super(MSCB, self).__init__() |
| 119 | |
| 120 | self.in_channels = in_channels |
| 121 | self.out_channels = out_channels |
| 122 | self.stride = stride |
| 123 | self.kernel_sizes = kernel_sizes |
| 124 | self.expansion_factor = expansion_factor |
| 125 | self.dw_parallel = dw_parallel |
| 126 | self.add = add |
| 127 | self.activation = activation |
| 128 | self.n_scales = len(self.kernel_sizes) |
| 129 | # check stride value |
| 130 | assert self.stride in [1, 2] |
| 131 | # Skip connection if stride is 1 |
| 132 | self.use_skip_connection = True if self.stride == 1 else False |
| 133 | |
| 134 | # expansion factor |
| 135 | self.ex_channels = int(self.in_channels * self.expansion_factor) |
| 136 | self.pconv1 = nn.Sequential( |
| 137 | # pointwise convolution |
| 138 | nn.Conv2d(self.in_channels, self.ex_channels, 1, 1, 0, bias=False), |
| 139 | nn.BatchNorm2d(self.ex_channels), |
| 140 | act_layer(self.activation, inplace=True) |
| 141 | ) |
| 142 | self.msdc = MSDC(self.ex_channels, self.kernel_sizes, self.stride, self.activation, dw_parallel=self.dw_parallel) |
| 143 | if self.add == True: |
| 144 | self.combined_channels = self.ex_channels*1 |
| 145 | else: |
| 146 | self.combined_channels = self.ex_channels*self.n_scales |
| 147 | self.pconv2 = nn.Sequential( |
| 148 | # pointwise convolution |
| 149 | nn.Conv2d(self.combined_channels, self.out_channels, 1, 1, 0, bias=False), |
| 150 | nn.BatchNorm2d(self.out_channels), |
| 151 | ) |
| 152 | if self.use_skip_connection and (self.in_channels != self.out_channels): |
| 153 | self.conv1x1 = nn.Conv2d(self.in_channels, self.out_channels, 1, 1, 0, bias=False) |
| 154 | self.init_weights('normal') |
| 155 | |
| 156 | def init_weights(self, scheme=''): |
| 157 | named_apply(partial(_init_weights, scheme=scheme), self) |
| 158 | |
| 159 | def forward(self, x): |
| 160 | pout1 = self.pconv1(x) |
| 161 | msdc_outs = self.msdc(pout1) |
| 162 | if self.add == True: |
| 163 | dout = 0 |
| 164 | for dwout in msdc_outs: |
| 165 | dout = dout + dwout |
| 166 | else: |
| 167 | dout = torch.cat(msdc_outs, dim=1) |
| 168 | dout = channel_shuffle(dout, gcd(self.combined_channels,self.out_channels)) |
| 169 | out = self.pconv2(dout) |
| 170 | if self.use_skip_connection: |