PyTorch nn.Transformer 用法笔记

本文项目链接: Transformer chatbot

nn.Transformer 结构:

抽象层1: nn.Transformer (不常用，不是特别灵活)
抽象层2: nn.TransformerEncoder, nn.TransformerDecoder (常用，比较灵活)
抽象层3: nn.TransformerEncoderLayer, nn.TransformerDecoderLayer (常用于自定义架构, 非常灵活)

这里直接忽略 nn.Transformer 的使用方法，只讲解第2, 3层

单层 transformer

`nn.TransformerEncoderLayer` 参数和期望输入输出维度

输入输出维度一致: pyTorch保证每个Layer的输入输出维度都是一致的，其实也很好理解:

如果不看LayerNorm层的话, TransformerEncoderLayer只有 attention layer 和 FNN两个组成
attention layer输入输出维度是一定是一样的，这个可以证明 (输入 d_model 维度向量，输出 d_model 维度向量)
而 FNN 接到之后为了保证 残差连接正常 也要保证输入输出维度是一样的
所以 nn.TransformerEncoderLayer 输入输出维度一样，如果想改变的话请自行添加新的FNN层

固定FNN结构: 在pyTorch中，feedforward的结构被固定为两层FNN，一层为隐藏层（relu激活），一层为线性输出层（符合输出维度），为什么固定结构呢？因为这是大量实验的经验之谈，可以不用改这个。FNN的结构如下:

class feedforward(nn.Module):
    def __init__ (self, d_model, dim_feedforward = 2048): # 隐藏层维度默认 2048
        '''
        维度变化:
        d_model - dim_feedforward - d_model
        '''
        super().__init__()
        self.hidden_layer = nn.Linear(d_model, dim_feedforward)
        self.out_layer = nn.Linear(dim_feedforward, d_model) # 输入输出维度一致

intial 参数

class TransformerEncoderLayer(
    d_model : int, #输入与输出的sequence中每个元素的维度 (注意，输入输出维度要一样，主要是为了叠多层时不要出乱子)
    nhead : int, #多头注意力头数，应该要整除 d_model, 因为多头注意力实现就是把维度平均分给每一个注意力头并行处理（和GAT中的多头注意力不太一样）
    dim_feedforward : int, # 这是feed forward隐藏层的维度，
    dropout : float, # dropout 默认 0.1
    activation : str, # FNN隐藏层用的激活函数默认 relu
    layer_norm_eps : float, # layer norm用的参数
    batch_first : bool, # 默认 false，建议改成True，因为习惯上在NLP中我们把batch维度放在第二维，seq放到第一维，而设定这个为 True 就和CNN的习惯统一了
    # 补充: True: (batch, seq, feature). False: False (seq, batch, feature)
    norm_first : bool, # True: norm - attention - norm - FNN, False: attention - norm - FNN - norm, 默认 False
    bias : bool, # FNN是否要bias？ 默认 True
    device, # 设备
    dtype # 格式
)

Encoder输入输出维度（batch_first = True）:

输入和输出都是 [Batch, seqlen, d_model]，其中encoder输出我们一般叫做 memory
有人可能会问：Encoder不是要负责输出 key 和 value 吗？怎么每个向量只对应一个输出值呢？
实际上在Transformer中 Encoder 确实只输出 [Batch, seqlen, d_model] 的输出向量，只不过我们会在 Decoder 中用两个linear层把 Encoder 的单个输出转换为 key和value，所以 我们把转换任务交给了Decoder，Encoder不会管这个！
另外，输入维度保持 [Batch, seqlen, d_model] 方便叠加多层 nn.TransformerEncoderLayer, 这样在Encoder层与层之间维度并不会变化。当然也方便了 BERT 的实现
forward() 函数参数:

def forward(
    self,
    src : Tensor,
    src_mask,
    src_key_padding_mask,
    is_causal : False # 因果注意力，如果为true，自动生成 look-ahead 来防止看到未来的东西
)

Encoder输入 mask

mask 实现原理是直接加到注意力矩阵上，使得注意力矩阵中被mask的地方恒为-inf而失效
- 注意力矩阵的维度是 [B, nhead, seqlen, seqlen], 其中最后两个维度分别是 [query, key]， $A_{ij}$ 表示第i个query对第j个key的注意力分数。
- 为什么是 -inf？：因为所有的mask都是在softmax之前加到注意力矩阵上的，所以softmax之后 -inf 就是0了
mask的工作原理非常简单，可以直接把mask 加 (+) 到输入tensor上，所以注意mask的样例： 0 表示该位置正常, -inf 表示该位置要被mask

Encoder 支持两种mask：`src_mask`, `padding_mask`

src_mask 用于人为屏蔽某些输入，虽然没有decoder掩盖前向注意力的需求，但是某些情况下我们可能希望mask一些数据，作为保留
- src_mask 维度是 [seqlen, seqlen], unsqueeze两个维度（batch，注意力）后正好可广播为 自注意力矩阵 维度（[B, nhead, seqlen, seqlen]），实践中我们也是直接把mask叠加到注意力矩阵上的 (在softmax前)
padding_mask 用于屏蔽padding的值，因为同一个batch训练集中的seq长度可能不一样，所以需要padding填充，这个是为了防止填充影响结果的mask，格式和src_mask 一样
- padding_mask维度是 [B, seqlen]，我们unsqueeze一个注意力头维度: [B, 1, seqlen],
- 然后 unsqueeze 一个 query 维度 [B, 1, 1, seqlen]
- 对最后一个维度(key维度) copy seqlen 次 这个可以通过pyTorch的广播机制实现，直接加就行了
- 这样就实现了 禁止任何被 mask 的值对应的 key 能被 query 注意到，i.e. 禁止了mask值被注意，从而防止对结果产生影响
- 注意，Transformer的 padding mask 只是禁止被mask值对应的key被query注意到，但是不妨碍被mask值本身作为query去注意其他值！！！
- 所以注意，被mask的值还是会提供一个输出向量，因为输出向量的生成只和 value 和 query 有关。mask 只能禁止 mask 值参与其他有用的 query 对应的输出计算，但是不能阻止他输出一个value，而这个value其实是我们不想要的，所以我们要屏蔽掉他
- 由encoder本事没法处理value的mask，所以我们交给解码器，解码器有一个 memory_mask (对应 src_mask) 和 memory_key_padding_mask (对应src_key_padding_mask) 就是干这个的，如果你要直接接 loss function，注意你也要在loss中应用mask！！！
- 可以看见nn.transformer的命名也非常准确，memory_key_padding_mask 只 mask key，不 mask query 和 value!, 输出值的mask靠的是Decoder或者损失函数！！！

`nn.TransformerDecoderLayer` 参数和期望输入输出维度

Decoder 相对于 Encoder来讲最大的区别就是他有两个attention，一个是 look-ahead attention，一个是 cross attention

`nn.TransformerDecoderLayer` 初始化参数和encoder相同:

class TransformerEncoderLayer(
    d_model : int, #输入与输出的sequence中每个元素的维度 (注意，输入输出维度要一样，主要是为了叠多层时不要出乱子)
    nhead : int, #多头注意力头数，应该要整除 d_model, 因为多头注意力实现就是把维度平均分给每一个注意力头并行处理（和GAT中的多头注意力不太一样）
    dim_feedforward : int, # 这是feed forward隐藏层的维度，
    dropout : float, # dropout 默认 0.1
    activation : str, # FNN隐藏层用的激活函数默认 relu
    layer_norm_eps : float, # layer norm用的参数
    batch_first : bool, # 默认 false，建议改成True，因为习惯上在NLP中我们把batch维度放在第二维，seq放到第一维，而设定这个为 True 就和CNN的习惯统一了
    # 补充: True: (batch, seq, feature). False: False (seq, batch, feature)
    norm_first : bool, # True: norm - attention - norm - FNN, False: attention - norm - FNN - norm, 默认 False
    bias : bool, # FNN是否要bias？ 默认 True
    device, # 设备
    dtype # 格式
)

`nn.TransformerDecoderLayer` 输入输出范式:

decoder 需要有两个输入，一个是 tgt输入, 一个是 memory输入, 一个和输入维度相同的输出，Decoder本身不包括 auto regression，需要自行处理
- tgt输入: 维度为 [B, seqlen2, d_model]
- memory: 维度为 [B, seqlen, d_model], （注意，一般为了简化， encoder-decoder的 d_model 是会保持一样的，当然理论上可以不一样）
- 注意，memory 就是 encoder 的输出，encoder输出的 seqlen 和 decoder输出的 seqlen2 是可以不一样的，也就是说交叉注意力矩阵可能不是方阵！而是 [seqlen2, seqlen] (query, key)
- 输出：为每个输入的向量创建一个对应的输出向量，维度不变，即: [B, seqlen2, d_model]
- 补: 在 seq2seq任务时，每个位置的向量对应的输出的向量应该表示的是: 预测的下一个向量是什么，但是由于 decoder 并不会采用自回归，而是 对每个输入都预测下一个输出，即 teacher forcing
- 注意，为了方便 teacher forcing 训练，Decoder 不提供自回归推理方法，需要自己创建一个新的method来推理
- 所以构建auto regression训练集时注意了，decoder的训练集的 (x, y) 中 y 应该仅仅是 x 向左 shift 一格，达成 “预测下一步输入” 的目的

mask

由于要处理两个attention层，所以 decoder 提供了四个mask接口
1. 用于lood-ahead attention 层：
- tgt_mask: 用于处理 lood-ahead mask
- tgt_key_padding_mask: 用于处理 tgt 输入的padding
1. 用于cross attention 层：
- 由于 encoder 的mask只能遮盖 key，不能遮盖 query 和 value，而这两个值不能参与attention运算。而 decoder 接受了 encoder 的 value 后创建了对应的 key 和 value，所以此时的 encoder 的 mask 没法阻止被 mask 的值参与运算，所以我们要把 encoder 的 mask 也传给decoder，用于消除影响
- memory_mask: 对应 encoder 的 src_mask
- memory_key_padding_mask: 对应 encoder 的 src_key_padding_mask

其他 forward 参数

def forward (
    self,
    tgt,
    memoty,
    tgt_mask,
    memory_mask,
    tgt_key_padding_mask,
    memory_key_padding_mask,
    tgt_is_causal, # 自动在look-ahead attention layer 中创建 因果mask
    memory_is_casual # 自动在 cross attention layer 中创建 因果mask
)

多层 transformer

`nn.TransformerEncoder`, `nn.TransformerDecoder` 只不过是把单层的 encoder-decoder 叠成多层罢了:

如果直接使用这两个层的话，多层encoder只会使用最后一层encoder作为 memory，输入给所有 decoder layer
两者的forward方法和单层encoder-decoder完全一致，只有initial有一点区别:


class TransformerEncoder (
    encoder_layer : nn.TransformerEncoderLayer, # 单层 encoder
    num_layer : int, # 暴力叠加多层
    norm : object, # 多层encoder的最后在加一个norm层，可选
    enable_nested_tensor : bool, # 默认 True，在内部把 tensor 转化为 nested tensor
    mask_check : bool # 默认 True
)

class TransformerDecoder (
    decoder_layer : nn.TransformerEncoderLayer, # 单层 decoder
    num_layer : int, # 暴力叠加多层
    norm : object # 多层decoder的最后在加一个norm层，可选
)

推荐参考项目: 使用 nn.TransformerEncoder 与 nn.TransformerDecoder 实现简单 cahtbot

项目链接: Transformer chatbot
我将在接下来的章节中介绍我在项目中获得的经验。

一点 nn.Transformer 的训练心得:

padding mask 和 look-ahead mask 一个不能少
注意 nn.Transformer 输出本身就满足 teacher forcing，输入输出 sequence 长度都是一样的
注意loss function也要mask padding, 因为 padding mask 的机制，所以我们只防止了其他position注意到padding的position，没有禁止 padding 的 position 注意其他 position
batch 是一个非常有用的防止局部最优解的方法，建议最好就是 32 ~ 128 （我选择 64）
Dropout层除了输出层外几乎都可以加（比如 FNN 解嵌入层的输入就可以加上Dropout）
如果 batch 选太大了，loss是降不下来的
要善用 lr_schedular, 这个的超参非常难调，但是非常重要！
现代实现中常常用 nn.embedding 来进行位置编码，而不是传统的三角函数方法
transformer层堆叠的层数要和数据集规模匹配，否则叠再多也没什么用
transformer输出解嵌入其实就一个nn.Linear就可以了，最多两层，不要再多了，大头留给 Transformer 层
如无必要，不要改变隐藏层维度！
如果 pyTorch 无缘无故报什么 replacement error，那么大概率是你用了一些 Batch size 敏感的操作，然后最后一个 batch 不能整除，size小了一些导致你的 reshape 函数中自动调用了某些 replace 操作
在ChatBot中，也许Embedding维度比叠加Transformer层维度还要重要！！！ 1024 维就是会比 512 维学到更多东西！
在ChatBot中，优化器和学习率调度非常重要，本项目使用 hugging face 在 transformers 库中推荐使用的 AdamW + Warm up scheduler
即使有 scheduler, 最大学习率也不要调的太高，0.0001 数量级附近比较合适

nn.Transformer 结构:

单层 transformer

nn.TransformerEncoderLayer 参数和期望输入输出维度

输入输出维度一致: pyTorch保证每个Layer的输入输出维度都是一致的，其实也很好理解:

固定FNN结构: 在pyTorch中，feedforward的结构被固定为 两层FNN，一层为隐藏层（relu激活），一层为线性输出层（符合输出维度），为什么固定结构呢？因为这是大量实验的经验之谈，可以不用改这个。FNN的结构如下:

intial 参数

Encoder输入输出维度 （batch_first = True）:

Encoder输入 mask

Encoder 支持两种mask：src_mask, padding_mask

nn.TransformerDecoderLayer 参数和期望输入输出维度

Decoder 相对于 Encoder来讲最大的区别就是 他有两个attention，一个是 look-ahead attention，一个是 cross attention

nn.TransformerDecoderLayer 初始化参数和encoder相同:

nn.TransformerDecoderLayer 输入输出范式:

mask

其他 forward 参数

多层 transformer

nn.TransformerEncoder, nn.TransformerDecoder 只不过是把 单层的 encoder-decoder 叠成多层罢了:

推荐参考项目: 使用 nn.TransformerEncoder 与 nn.TransformerDecoder 实现简单 cahtbot

一点 nn.Transformer 的训练心得:

`nn.TransformerEncoderLayer` 参数和期望输入输出维度

固定FNN结构: 在pyTorch中，feedforward的结构被固定为两层FNN，一层为隐藏层（relu激活），一层为线性输出层（符合输出维度），为什么固定结构呢？因为这是大量实验的经验之谈，可以不用改这个。FNN的结构如下:

Encoder输入输出维度（batch_first = True）:

Encoder 支持两种mask：`src_mask`, `padding_mask`

`nn.TransformerDecoderLayer` 参数和期望输入输出维度

Decoder 相对于 Encoder来讲最大的区别就是他有两个attention，一个是 look-ahead attention，一个是 cross attention

`nn.TransformerDecoderLayer` 初始化参数和encoder相同:

`nn.TransformerDecoderLayer` 输入输出范式:

`nn.TransformerEncoder`, `nn.TransformerDecoder` 只不过是把单层的 encoder-decoder 叠成多层罢了: