| 3 | |
| 4 | |
| 5 | class BlockA(nn.Module): |
| 6 | def __init__(self, in_channels=64, out_channels=64, inter_channels=64, mlp_ratio=4.): |
| 7 | super(BlockA, self).__init__() |
| 8 | inter_channels = in_channels |
| 9 | self.conv = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) |
| 10 | self.norm1 = nn.LayerNorm(inter_channels) |
| 11 | self.ffn = MLPLayer(in_features=inter_channels, |
| 12 | hidden_features=int(inter_channels * mlp_ratio), |
| 13 | act_layer=nn.GELU, |
| 14 | drop=0.) |
| 15 | self.norm2 = nn.LayerNorm(inter_channels) |
| 16 | |
| 17 | def forward(self, x): |
| 18 | B, C, H, W = x.shape |
| 19 | _x = self.conv(x) |
| 20 | _x = _x.flatten(2).transpose(1, 2) |
| 21 | _x = self.norm1(_x) |
| 22 | x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
| 23 | |
| 24 | x = x + _x |
| 25 | _x1 = self.ffn(x) |
| 26 | _x1 = self.norm2(_x1) |
| 27 | _x1 = _x1.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
| 28 | x = x + _x1 |
| 29 | return x |
nothing calls this directly
no outgoing calls
no test coverage detected