| 141 | |
| 142 | # Stage 2 |
| 143 | class ChannelEmbed(nn.Module): |
| 144 | def __init__(self, in_channels, out_channels, reduction=1, norm_layer=nn.BatchNorm2d): |
| 145 | super(ChannelEmbed, self).__init__() |
| 146 | self.out_channels = out_channels |
| 147 | self.residual = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) |
| 148 | self.channel_embed = nn.Sequential( |
| 149 | nn.Conv2d(in_channels, out_channels // reduction, kernel_size=1, bias=True), |
| 150 | nn.Conv2d( |
| 151 | out_channels // reduction, |
| 152 | out_channels // reduction, |
| 153 | kernel_size=3, |
| 154 | stride=1, |
| 155 | padding=1, |
| 156 | bias=True, |
| 157 | groups=out_channels // reduction, |
| 158 | ), |
| 159 | nn.ReLU(inplace=True), |
| 160 | nn.Conv2d(out_channels // reduction, out_channels, kernel_size=1, bias=True), |
| 161 | norm_layer(out_channels), |
| 162 | ) |
| 163 | self.norm = norm_layer(out_channels) |
| 164 | |
| 165 | def forward(self, x, H, W): |
| 166 | B, N, _C = x.shape |
| 167 | x = x.permute(0, 2, 1).reshape(B, _C, H, W).contiguous() |
| 168 | residual = self.residual(x) |
| 169 | x = self.channel_embed(x) |
| 170 | out = self.norm(residual + x) |
| 171 | return out |
| 172 | |
| 173 | |
| 174 | # FFM |