| import torch.nn as nn | |
| from models.modules.mlp import MLPLayer | |
| class BlockA(nn.Module): | |
| def __init__(self, in_channels=64, out_channels=64, inter_channels=64, mlp_ratio=4.): | |
| super(BlockA, self).__init__() | |
| inter_channels = in_channels | |
| self.conv = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) | |
| self.norm1 = nn.LayerNorm(inter_channels) | |
| self.ffn = MLPLayer(in_features=inter_channels, | |
| hidden_features=int(inter_channels * mlp_ratio), | |
| act_layer=nn.GELU, | |
| drop=0.) | |
| self.norm2 = nn.LayerNorm(inter_channels) | |
| def forward(self, x): | |
| B, C, H, W = x.shape | |
| _x = self.conv(x) | |
| _x = _x.flatten(2).transpose(1, 2) | |
| _x = self.norm1(_x) | |
| x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() | |
| x = x + _x | |
| _x1 = self.ffn(x) | |
| _x1 = self.norm2(_x1) | |
| _x1 = _x1.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() | |
| x = x + _x1 | |
| return x |