Transformer模型代码逐行解释，免费附上所有代码

1. 第一段代码

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

1. class PreNorm(nn.Module):: 这行定义了一个类 PreNorm，它继承自 nn.Module，意味着它是一个 PyTorch 模块。

2. def __init__(self, dim, fn):: 这是 PreNorm 类的初始化函数。它接受两个参数：dim 和 fn。dim 是输入数据的维度，fn 是一个函数或者模块，将应用于数据。

3. super().__init__(): 这行调用了父类 nn.Module 的初始化函数，确保正确地初始化模块。

4. self.norm = nn.LayerNorm(dim): 这行创建了一个 Layer Normalization 层，并将其赋值给 self.norm。Layer Normalization 是一种归一化技术，用于调整输入数据的分布，使其更易于处理。

5. self.fn = fn: 这行将传入的函数或模块赋值给 self.fn，以备稍后在前向传播时使用。

6. def forward(self, x, **kwargs):: 这是 PreNorm 类的前向传播函数。它接受输入 x 和可选的关键字参数 kwargs。

7. return self.fn(self.norm(x), **kwargs): 这行调用了 self.norm 对输入 x 进行归一化，然后将归一化后的结果传递给 self.fn 进行处理。这就是所谓的 “Pre-Norm”，即在应用函数或模块之前对输入进行归一化处理。

8. class FeedForward(nn.Module):: 这行定义了另一个类 FeedForward，同样是继承自 nn.Module。

9. def __init__(self, dim, hidden_dim, dropout=0.):: 这是 FeedForward 类的初始化函数。它接受三个参数：dim 是输入数据的维度，hidden_dim 是隐藏层的维度，dropout 是 dropout 的比率，默认为 0。

10. super().__init__(): 这行调用了父类 nn.Module 的初始化函数。

11. self.net = nn.Sequential(: 这行创建了一个 Sequential 容器，用于按顺序组织多个层或操作。

12. nn.Linear(dim, hidden_dim),: 这行添加了一个线性层，将输入维度为 dim 的数据线性映射到维度为 hidden_dim 的空间。

13. nn.GELU(),: 这行添加了一个 GELU 激活函数。GELU 是一种激活函数，它在近期得到了广泛的应用，特别是在 Transformer 模型中。

14. nn.Dropout(dropout),: 这行添加了一个 Dropout 层，以减少模型过拟合的风险。

15. nn.Linear(hidden_dim, dim),: 这行添加了另一个线性层，将隐藏层的输出映射回原始维度 dim。

16. nn.Dropout(dropout): 这行再次添加了一个 Dropout 层，对输出进行一定比例的随机丢弃。

17. def forward(self, x):: 这是 FeedForward 类的前向传播函数，它接受输入 x。

18. return self.net(x): 这行调用了之前定义的 Sequential 容器 self.net，将输入 x 传递给其中的每一层，实现了前向传播过程。

2.第二段代码

class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim=-1)
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        x = rearrange(x, 'b 1 n l -> b n l')
        #将x转变为qkv
        qkv = self.to_qkv(x).chunk(3, dim=-1)  # 对tensor进行分块

        q, k, v = \
            map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
        attn = self.attend(dots)

        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        out = rearrange(out, 'b n l -> b 1 n l')
        return self.to_out(out)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

这段代码定义了一个名为 Attention 的 PyTorch 模块，用于实现自注意力机制。

1. class Attention(nn.Module):: 这行定义了一个类 Attention，它继承自 nn.Module，意味着它是一个 PyTorch 模块。

2. def __init__(self, dim, heads=8, dim_head=64, dropout=0.):: 这是 Attention 类的初始化函数。它接受四个参数：dim 是输入数据的维度，heads 是注意力头的数量，默认为 8，dim_head 是每个注意力头的维度，默认为 64，dropout 是 dropout 的比率，默认为 0。

3. super().__init__(): 这行调用了父类 nn.Module 的初始化函数。

4. inner_dim = dim_head * heads: 这行计算了注意力头的总维度，即 dim_head 乘以 heads。

5. project_out = not (heads == 1 and dim_head == dim): 这行判断是否需要投影输出。如果 heads 等于 1 并且 dim_head 等于 dim，则不需要投影输出。

6. self.heads = heads: 这行将 heads 赋值给 self.heads，以备后续使用。

7. self.scale = dim_head ** -0.5: 这行计算了缩放因子，用于缩放点积注意力的分数，以避免梯度爆炸或消失。

8. self.attend = nn.Softmax(dim=-1): 这行创建了一个 Softmax 层，用于计算注意力权重。

9. self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False): 这行创建了一个线性层，用于将输入数据映射到 Query、Key 和 Value 矩阵。inner_dim * 3 表示输出的特征维度是 inner_dim 的三倍，因为需要分别计算 Q、K 和 V。

10. self.to_out = nn.Sequential(: 这行创建了一个 Sequential 容器，用于按顺序组织多个层或操作。如果 project_out 为真，则执行投影输出的操作，否则使用恒等映射。

11. nn.Linear(inner_dim, dim),: 这行创建了一个线性层，将内部维度 inner_dim 的数据映射回原始维度 dim。

12. nn.Dropout(dropout): 这行添加了一个 Dropout 层，对输出进行一定比例的随机丢弃。

13. if project_out else nn.Identity(): 这是一个三元表达式，如果 project_out 为真，则执行前面的操作，否则返回 nn.Identity()，即恒等映射。

14. def forward(self, x):: 这是 Attention 类的前向传播函数，它接受输入 x。

15. x = rearrange(x, 'b 1 n l -> b n l'): 这行将输入张量的维度重新排列，将大小为 1 的维度（通常用于序列批次维度）移动到中间。

16. qkv = self.to_qkv(x).chunk(3, dim=-1): 这行调用 self.to_qkv 将输入 x 映射为 Query、Key 和 Value，然后使用 chunk 函数将结果分割成 3 个张量。

17. map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv): 这行使用 map 函数对分割后的 Query、Key 和 Value 张量进行重排列，以适应多头自注意力的计算需求。

18. dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale: 这行计算了 Query 和 Key 的点积注意力分数，并乘以缩放因子。

19. attn = self.attend(dots): 这行通过 Softmax 函数计算了点积注意力的权重。

20. out = torch.matmul(attn, v): 这行计算了加权的 Value，得到了注意力输出。

21. out = rearrange(out, 'b h n d -> b n (h d)'): 这行将输出重排列为原始形状，以便后续处理。

22. out = rearrange(out, 'b n l -> b 1 n l'): 这行将输出张量的维度再次重新排列，将大小为 1 的维度移回到最前面。

23. return self.to_out(out): 这行将输出传递给 self.to_out 进行投影输出或恒等映射，并返回最终结果。

3. 第三段代码

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.pooling = nn.AvgPool2d(kernel_size=(1, 2), stride=(1, 2))
        self.layers = nn.ModuleList([])
        for i in range(depth):
            self.layers.append(nn.ModuleList([
                # PreNorm(dim, ResNetBlock()),
                # PreNorm(dim//(2**i), ConvAttention(dim//(2**i), heads, dim_head//(2**i), dropout)),
                # PreNorm(dim//(2**(i+1)), FeedForward(dim//(2**(i+1)), mlp_dim, dropout))
                PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
            ]))

    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x

        return x
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

这段代码定义了一个名为 Transformer 的 PyTorch 模块，实现了一个简化的 Transformer

1. class Transformer(nn.Module):: 这行定义了一个类 Transformer，它继承自 nn.Module，意味着它是一个 PyTorch 模块。

2. def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):: 这是 Transformer 类的初始化函数。它接受六个参数：dim 是输入数据的维度，depth 是 Transformer 的层数，heads 是注意力头的数量，dim_head 是每个注意力头的维度，mlp_dim 是 MLP 隐藏层的维度，dropout 是 dropout 的比率，默认为 0。

3. super().__init__(): 这行调用了父类 nn.Module 的初始化函数。

4. self.pooling = nn.AvgPool2d(kernel_size=(1, 2), stride=(1, 2)): 这行创建了一个平均池化层，用于将输入的时间序列数据降采样。

5. self.layers = nn.ModuleList([]): 这行创建了一个空的 nn.ModuleList，用于存储 Transformer 的每一层。

6. for i in range(depth):: 这是一个 for 循环，用于构建 Transformer 的每一层。

7. self.layers.append(nn.ModuleList([: 这行向 self.layers 中添加一个新的模块列表，用于存储当前层的注意力机制和前馈网络。

8. PreNorm(dim, Attention(dim, heads, dim_head, dropout)),: 这行创建了一个 PreNorm 模块，其输入是 dim 维度的数据，应用了自注意力机制（通过 Attention 类实现），其中 heads 是注意力头的数量，dim_head 是每个注意力头的维度。

9. PreNorm(dim, FeedForward(dim, mlp_dim, dropout)): 这行创建了一个 PreNorm 模块，其输入也是 dim 维度的数据，应用了前馈网络（通过 FeedForward 类实现），其中 mlp_dim 是 MLP 隐藏层的维度。

10. def forward(self, x):: 这是 Transformer 类的前向传播函数，它接受输入 x。

11. for attn, ff in self.layers:: 这是一个 for 循环，遍历每一层的注意力机制和前馈网络。

12. x = attn(x) + x: 这行将输入 x 传递给注意力机制模块 attn，然后将注意力机制的输出与原始输入相加，实现了残差连接。

13. x = ff(x) + x: 这行将残差连接后的输出 x 传递给前馈网络模块 ff，然后将前馈网络的输出与之前的输出再次相加，实现了残差连接。

14. return x: 这行返回最终的输出 x。

4. 第四段代码

class TimeTransformer(nn.Module):
    def __init__(self, *, input_dim,  num_patches=16, dim, depth, heads, mlp_dim,
                 pool='cls', channels=1, dim_head, emb_dropout=0., dropout=0.):
        super(TimeTransformer, self).__init__()

        # self.to_patch_embedding = Embedding(input_dim, dim)
        self.to_patch_embedding = self.to_patch_embedding = nn.Sequential(
            Rearrange('b 1 (n d) -> b 1 n d', n=num_patches),
            nn.Linear(input_dim//num_patches, dim)
        )
        self.cls_token = nn.Parameter(torch.randn(1, 1, 1, dim))  # [1, 1, 1, dim] 随机数
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))  # [1, num_patches+1, dim] 随机数
        self.dropout = nn.Dropout(emb_dropout)

        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
        self.pool = pool
        self.to_latent = nn.Identity()  # 这个恒等函数，就如同名字占位符，并没有实际操作

    def forward(self, rawdata):
        TimeSignals = rawdata   # Get Time Domain Signals
        TimeSignals = rearrange(TimeSignals, 'b l -> b 1 l')
        # print(TimeSignals.shape, rawdata.shape)

        x = self.to_patch_embedding(TimeSignals)
        b, _, n, _ = x.shape      # x: [batch_size, channels, num_patches, dim]

        cls_tokens = repeat(self.cls_token, '() c n d -> b c n d', b=b)  # cls_tokens: [batch_size, c, num_patches, dim]
        x = torch.cat((cls_tokens, x), dim=2)  # x: [batch_size, c, num_patches+1, dim]
        # print(x.shape)
        #x += self.pos_embedding[:, :(n + 1)]  # 添加位置编码：x: [batch_size, c, num_patches+1, dim]
        x = self.dropout(x)

        x = self.transformer(x)     # x: [batch_size, c, num_patches+1, dim]
        x = x.mean(dim=1) if self.pool == 'mean' else x[:, :, 0]  # x: [batch_size, c, 1, dim]
        x = self.to_latent(x)
        return x
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

这段代码定义了一个名为 TimeTransformer 的 PyTorch 模块，用于处理时间序列数据的 Transformer 架构。

1. class TimeTransformer(nn.Module):: 这行定义了一个类 TimeTransformer，它继承自 nn.Module，意味着它是一个 PyTorch 模块。

2. def __init__(self, *, input_dim, num_patches=16, dim, depth, heads, mlp_dim, pool='cls', channels=1, dim_head, emb_dropout=0., dropout=0.):: 这是 TimeTransformer 类的初始化函数。它接受多个命名关键字参数：input_dim 是输入数据的维度，num_patches 是图像被分成的块数，默认为 16，dim 是 Transformer 模型的特征维度，depth 是 Transformer 模型的层数，heads 是注意力头的数量，mlp_dim 是 MLP 隐藏层的维度，pool 是池化方式，默认为 ‘cls’，channels 是输入数据的通道数，默认为 1，dim_head 是每个注意力头的维度，emb_dropout 是输入嵌入的 dropout 比率，默认为 0，dropout 是 Transformer 模型中的 dropout 比率，默认为 0。

3. super(TimeTransformer, self).__init__(): 这行调用了父类 nn.Module 的初始化函数。

4. self.to_patch_embedding = self.to_patch_embedding = nn.Sequential(: 这行定义了将原始时间序列数据转换为补丁嵌入的过程。首先使用 Rearrange 操作将数据重排为合适的形状，然后通过线性层将其映射到特征维度 dim。

5. self.cls_token = nn.Parameter(torch.randn(1, 1, 1, dim)): 这行定义了一个可学习的类标记向量，它的维度是 [1, 1, 1, dim]，初始值为随机数。

6. self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim)): 这行定义了一个可学习的位置编码矩阵，其维度为 [1, num_patches+1, dim]，初始值为随机数。这个位置编码将被加到补丁嵌入中，以表示每个补丁的位置信息。

7. self.dropout = nn.Dropout(emb_dropout): 这行创建了一个 Dropout 层，用于输入嵌入的 dropout 操作。

8. self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout): 这行创建了一个 Transformer 模型，其中包含自注意力机制和前馈网络。

9. self.pool = pool: 这行将池化方式保存在实例中，以备后续使用。

10. self.to_latent = nn.Identity(): 这行创建了一个恒等映射，作为模型输出的最后一层，以便根据需要添加额外的后处理步骤。

11. def forward(self, rawdata):: 这是 TimeTransformer 类的前向传播函数，它接受原始时间序列数据 rawdata。

12. TimeSignals = rearrange(rawdata, 'b l -> b 1 l'): 这行使用 rearrange 函数将原始时间序列数据重排为合适的形状，以适应输入嵌入的处理方式。

13. x = self.to_patch_embedding(TimeSignals): 这行将重排后的数据传递给补丁嵌入层进行处理。

14. cls_tokens = repeat(self.cls_token, '() c n d -> b c n d', b=b): 这行使用 repeat 函数将类标记向量复制成与补丁嵌入相同的形状，并添加到补丁嵌入中。

15. x = torch.cat((cls_tokens, x), dim=2): 这行将类标记向量和补丁嵌入沿着第二个维度进行拼接，形成模型的输入。

16. x = self.dropout(x): 这行对输入进行 dropout 操作。

17. x = self.transformer(x): 这行将输入传递给 Transformer 模型进行处理，得到模型的输出。

18. x = x.mean(dim=1) if self.pool == 'mean' else x[:, :, 0]: 这行根据指定的池化方式对模型输出进行池化操作，如果 self.pool 为 ‘mean’，则取平均值，否则只取类标记的输出。

19. x = self.to_latent(x): 这行通过恒等映射将模型输出传递给最后的输出层。

20. return x: 这行返回最终的输出。