29 lines
1.0 KiB
Python
29 lines
1.0 KiB
Python
import torch.nn as nn
|
|
from BiRefNet.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 |