Next Convolution Block
| 171 | |
| 172 | |
| 173 | class NCB(nn.Module): |
| 174 | """ |
| 175 | Next Convolution Block |
| 176 | """ |
| 177 | def __init__(self, in_channels, out_channels, stride=1, path_dropout=0, |
| 178 | drop=0, head_dim=32, mlp_ratio=3): |
| 179 | super(NCB, self).__init__() |
| 180 | self.in_channels = in_channels |
| 181 | self.out_channels = out_channels |
| 182 | norm_layer = partial(nn.BatchNorm2d, eps=NORM_EPS) |
| 183 | assert out_channels % head_dim == 0 |
| 184 | |
| 185 | self.patch_embed = PatchEmbed(in_channels, out_channels, stride) |
| 186 | self.mhca = MHCA(out_channels, head_dim) |
| 187 | self.attention_path_dropout = DropPath(path_dropout) |
| 188 | |
| 189 | self.norm = norm_layer(out_channels) |
| 190 | self.mlp = Mlp(out_channels, mlp_ratio=mlp_ratio, drop=drop, bias=True) |
| 191 | self.mlp_path_dropout = DropPath(path_dropout) |
| 192 | self.is_bn_merged = False |
| 193 | |
| 194 | def merge_bn(self): |
| 195 | if not self.is_bn_merged: |
| 196 | self.mlp.merge_bn(self.norm) |
| 197 | self.is_bn_merged = True |
| 198 | |
| 199 | def forward(self, x): |
| 200 | x = self.patch_embed(x) |
| 201 | x = x + self.attention_path_dropout(self.mhca(x)) |
| 202 | if not torch.onnx.is_in_onnx_export() and not self.is_bn_merged: |
| 203 | out = self.norm(x) |
| 204 | else: |
| 205 | out = x |
| 206 | x = x + self.mlp_path_dropout(self.mlp(out)) |
| 207 | return x |
| 208 | |
| 209 | |
| 210 | class E_MHSA(nn.Module): |