如何将归一化模块和别的模块写在一起

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   +     前馈网络