class PreNorm(nn.Module):
def __init__(self, dim, fn):
super().__init__()
self.fn = fn
self.norm = nn.LayerNorm(dim)
def forward(self, x, *args, **kwargs):
x = self.norm(x)
return self.fn(x, *args, **kwargs)
注释:
这个模块先对输入进行LayerNorm归一化,然后再应用传入的函数
fn
。它可以被用作其他模块(例如前馈网络)之前的归一化步骤。
class FeedForward(nn.Module):
def __init__(self, dim, mult=4):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(dim, dim * mult, 1, 1, bias=False),# [b, dim, h, w] -> [b, dim * mult, h, w]
GELU(),
nn.Conv2d(dim * mult, dim * mult, 3, 1, 1,# [b, dim * mult, h, w] -> [b, dim * mult, h, w]
bias=False, groups=dim * mult),
GELU(),
nn.Conv2d(dim * mult, dim, 1, 1, bias=False), # [b, dim * mult, h, w] -> [b, dim, h, w]
)
注释:
这个
FeedForward
类实现了一个前馈神经网络,它包含了三个卷积层和两个 GELU 激活函数。
class IGAB(nn.Module):
def __init__(
self,
dim,
dim_head=64,
heads=8,
num_blocks=2,
):
super().__init__()
self.blocks = nn.ModuleList([])
for _ in range(num_blocks):
self.blocks.append(nn.ModuleList([
IG_MSA(dim=dim, dim_head=dim_head, heads=heads),# 注意力模块
PreNorm(dim, FeedForward(dim=dim))# 前馈网络
]))
注释:
PreNorm(dim, FeedForward(dim=dim))# LN + 前馈网络