跟李沐学AI：Transformer

Landy_Jay

已于 2024-10-15 20:06:27 修改

阅读量611

点赞数 17

文章标签：人工智能 transformer 深度学习

于 2024-10-12 18:51:44 首次发布

本文链接：https://blog.csdn.net/Landy_Jay/article/details/142793187

版权

Transformer架构

(图源:10.7. Transformer — 动手学深度学习 2.0.0 documentation)

基于编码器-解码器架构来处理序列对

与使用注意力的seq2seq不同，Transformer纯基于注意力

多头注意力（Multi-Head Attention）

(图源:10.5. 多头注意力 — 动手学深度学习 2.0.0 documentation)

对于同一QKV，我们希望抽取不同的信息。多头注意力机制相比于标准的注意力机制，允许模型捕捉一个序列的不同方面并对不同方面生成不同权重。

假设输入序列为： I love palying football.

假设embedding_size=3并且使用2个注意力头，序列可表示为一个(4, 3)的矩阵(sequence_length, embed_size)。

[ I ]         [ love ]       [ playing ]    [ football ]
[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [1.0, 1.1, 1.2]]  # Shape: (4, 3)

首先我们需要把embedding投影为Q，K，V。

权重矩阵的形状由输入特征数和头数共同决定，不同头的权重矩阵的形状是相同的。

$W_Q^{i} \epsilon R^{d_{feature}\times d_{head}}$

$W_K^{i} \epsilon R^{d_{feature}\times d_{head}}$

$W_V^{i} \epsilon R^{d_{feature}\times d_{head}}$

$d_{feature}$ 是输入的embedding_size，在本例中为3。

$d_{head}=d_{feature}/h$ 。其中h是注意力头的数量。

故本例中权重矩阵的形状为(3, 1)，其中1为向下取整得到。

对于头1：

Q1 = Wq1 * Embedding
K1 = Wk1 * Embedding
V1 = Wv1 * Embedding

假设头1的权重为：

$W_Q^{head1} =[[0.1],[0.2],[0.3]]$

$W_K^{head1} =[[0.4],[0.5],[0.6]]$

$W_V^{head1} =[[0.7],[0.8],[0.9]]$

那么我们可以列出输入embedding变换至Q的计算过程

$Q_I^{\text{head1}} = [0.1, 0.2, 0.3] * [[0.1], [0.2], [0.3]] = [0.14]$

$Q_{love}^{\text{head1}} = [0.32]$

$Q_{playing}^{\text{head1}} = [0.5]$

$Q_{football}^{\text{head1}} = [0.68]$

矩阵表示为：

$Q^{\text{head1}} = [[0.14], [0.32], [0.5], [0.68]]$

同理可得：

$K^{\text{head1}} = [[0.32], [0.77], [1.22], [1.67]]$

$V^{\text{head1}} = [[0.38], [0.92], [1.46], [2.0]]$

对于头2，我们需要不同的权重，计算过程相同

随后对于每个头，对Q，K进行点积并进行softmax计算注意力权重，将权重与V相乘计算加权和。本例中，QK点积的结果为一个(4, 4)的矩阵。将注意力权重与V相乘得到(4, 1)的注意力输出。

计算了每个头的注意力后，对输出进行拼接(concatenate)操作。再将拼接后的结果输入另一个线性层以得到最后的输出。本例中两个注意力头拼接后得到(4, 2)的输出。随后使用一个线性层将结果转换为需要的维度(4, 3)。

过程总结：

输入维度: (4, 3) (4 词元, 嵌入维度为 3).
每个头的Q、K、V： (4, 1) (4 tokens, 1 dimension per head).
每个头的输出：(4, 1).
每个头拼接后：(4, 2).
线性映射后的输出：(4, 3).

基于位置的前馈网络FFN

将输入形状由(b, n, d)变为(bn, d)

输出形状由(bn, d)变回(b, n, d)

FFN等价于两层kernel_size为1的一维卷积层

将输入由(b, n, d)变为(bn, d)有如下原因：

全连接层要求输入为2维，一维为batch_size，二维表示特征。
将(b, n, d)变为(bn, d)而不是(b, nd)是为了保留每个元素在序列中的位置信息（位置编码提供的位置信息）。
可以有效地将所有批次序列中的每个位置视为具有d特征的单个样本。这允许密集层对序列中的每个位置应用相同的权重集合。
同时保障了并行计算，(bn, d)可以并行处理多个序列，可以同时对整个序列进行操作，而不是按顺序处理每个元素，这大大加快了计算速度并减少了整体训练时间。

层归一化

什么是层归一化？

对每个样本的所有feature做归一化。意思是对每个样本独立地计算均值和方差。

假设输入 $x=[[1, 2], [3, 4], [5, 6]]$ 为(3, 2)的矩阵，样本归一化对每行计算均值和方差：

对第一个样本：均值为(1+2)/2=1.5，方差为 $(1-1.5)^2+(2-1.5)^2=0.25$ 。

对每行计算并归一化后得到输出 $\hat{x}=[[-1, 1], [-1, 1], [-1, 1]]$

对比BatchNorm

Batch normalization对输入的所有batch的每个feature批量计算均值和方差并进行归一化。对每个feature独立地计算均值和方差。

同样假设 $x=[[1, 2], [3, 4], [5, 6]]$ ，对于第一个特征计算均值和方差：

$mean=(1+3+5)/3=3$ ， $variance=(1-3)^2+(3-3)^3+(5-3)^2=4$

对每列计算并归一化后的到输出： $\hat{x}=[[-1, -1], [0, 0], [1, 1]]$

为什么是有layer norm而不是batch norm？

Transform的输入一般为一个序列，但是每个序列的有效长度(valid length)有所不同，会导致batch norm不稳定，故不适合batch norm。

对每个样本的元素做归一化，在长度变化时更加稳定。

信息传递

编码器中的输出 $y_1,\dots,y_n$ 作为解码器中第i个Transformer块中多头注意力的K和V，解码器的Q来自目标序列。

说明编码器和解码器中的Block个数和输出维度是相同的。

预测

预测时，解码器的掩码注意力和多头注意力的Q、K和V的来源分别如下：

Masked Attention:
- Q: 解码器在时间步 t 的状态（之前生成的tokens的表示）。
- K: 解码器在时间步 t 的状态（之前生成的tokens的表示）。
- V: 解码器在时间步 t 的状态（之前生成的tokens的表示）。
Multi-Head Attention（交叉注意力）:
- Q: 解码器在时间步 ttt 的状态。
- K: 编码器的最终输出（输入序列的上下文表示）。
- V: 编码器的最终输出（输入序列的上下文表示）。

在训练时，解码器的Masked Attention和Multi-Head Attention的Q、K和V分别由以下组成：

Masked Attention:
- Q: 解码器当前时间步的输入向量（包括所有已生成的tokens的表示）。
- K: 解码器当前时间步的输入向量（包括所有已生成的tokens的表示）。
- V: 解码器当前时间步的输入向量（包括所有已生成的tokens的表示）。
Multi-Head Attention（交叉注意力）:
- Q: 解码器当前时间步的输入向量（表示已生成的tokens）。
- K: 编码器的最终输出（输入序列的上下文表示）。
- V: 编码器的最终输出（输入序列的上下文表示）。

基于Pytorch的Transformer

MultiHeadAttention

class MultiHeadAttention(nn.Module):
    # num_hiddens: 输入的特征数
    def __init__(self, key_size, query_size, value_size, num_hiddens, num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        # heads的数量
        self.num_heads = num_heads
        # 点积注意力
        self.attention = d2l.DotProductAttention(dropout)
        # 计算Q的权重矩阵
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        # 计算K的权重矩阵
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        # 计算V的权重矩阵
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        # 对注意力输出进行变换的矩阵
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # QKV的shape：(batch_size, nums, num_hiddens)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads) 
        values = transpose_qkv(self.W_v(values), self.num_heads)
        
        # valid_lens的shape：(batch_size, ) or (batch_size, num_query)
        if valid_lens is not None:
            # 每个头都需要valid_lens的信息，因此需要复制num_heads次
            # repeats=self.num_heads, dim=0 表示沿着第一个维度（即 batch 维度）将 valid_lens 的每一项都重复 num_heads
            # 由于每个头都需要知道它正在处理的序列的有效长度，因此我们需要将 valid_lens 对应地复制给每一个头。
            valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)

        # output的shape：(batch_size*num_heads, num_query, num_hiddens/num_heads)
        # 将不同头的QKV合并成为一个大矩阵进行计算，提升计算并行度。即K=K1 concat K2 concat .. Kn
        output = self.attention(queries, keys, values, valid_lens)
        # output_concat的形状：(batch_size, num_query, num_hiddens)
        # 将不同头的结果拼接
        output_concat = transpose_output(output, self.num_heads)
        # 对注意力输出进行线性变换后输出
        return self.W_o(output_concat)

# 两个辅助函数
#@save
def transpose_qkv(X, num_heads):
    # 为了多注意头并行计算而变换形状
    # X的输入为：(batch_size, num_query, num_hiddens)
    # 输出为：(batch_szie, num_query, num_heads, num_hiddnes / num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)
    # 交换heads维和num_query维
    # 输出为：(batch_size, num_heads, num_query, num_hiddens / num_heads)
    X = X.permute(0, 2, 1, 3)
    # 将输出变为(batch_size * num_heads, num_query, num_hiddens / num_heads)用于并行计算
    return X.reshape(-1, X.shape[2], X.shape[3])

#@save
def transpose_output(X, num_heads):
    # 将合并的X还原num_heads维度
    # 输入为：(batch_size, num_heads, num_query, num_hiddens/num_heads)
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    # 转换为：(batch_size, num_query, num_heads, num_hiddens / num_heads)
    X = X.permute(0, 2, 1, 3)
    # 转换为原格式：(batch_size, num_query, num_hiddens)
    return X.reshape(X.shape[0], X.shape[1], -1)

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()

batch_size, num_queries = 2, 4
# valid_lens表示第一个batch的有效长度为3，第二个为2
num_kvpairs, valid_lens =  6, torch.tensor([3, 2])
'''
在编码器-解码器架构中，比如在Transformer模型里，解码器部分的每个时间步都会产生一个查询向量去与编码器的输出进行交互。
查询的数量通常是解码器当前时间步的数目，或者是整个目标序列的长度（如果一次性处理整个序列的话）。
对于自注意力机制（self-attention）(编码器中)，每个位置上的元素都作为一个查询去与其他所有元素交互，所以这里的查询数量就是输入序列的长度。
'''
X = torch.ones((batch_size, num_queries, num_hiddens))
'''
在编码器中的自注意力层，键-值对的数量等同于输入序列的长度。
在解码器中，键-值对可以来自于编码器的输出，这时它们的数量就是编码器输出序列的长度；或者来自解码器自身的先前输出，
这时它们的数量将是解码器已经产生的序列长度。
'''
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape
'''
Q与KV数量不相同，如何进行注意力计算？
Q.shape = (batch_size, num_queries, d_model) 
K.shape = (batch_size, num_kvpairs, d_model) 
Q @ K =  (batch_size, num_queries, num_kvpairs) = A
将结果进行softmax后再与V相乘
V.shape = (batch_size, num_kvpairs, d_model) 
A @ K = (batch_size, num_queries, d_model) 
'''

基于位置的前馈网络（FFN）

# 基于位置的前馈网络
class PositionWiseFFN(nn.Module):
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))

残差连接和层正则化

# 残差连接和规范化
#@save
class AddNorm(nn.Module):
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
        # normalized_shape参数应该是一个整数或整数元组，用来指定归一化操作所作用的维度大小。
        # 对于全连接层（Dense Layer），这通常是最后一个维度的大小，即特征的数量。
        self.ln = nn.LayerNorm(normalized_shape)

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)

位置编码

# 位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        # 初始化一个Dropout层，防止过拟合。
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的位置编码矩阵P，形状为(1, max_len, num_hiddens)。
        self.P = torch.zeros((1, max_len, num_hiddens))
        
        # 计算每个位置的编码值。X代表了位置索引除以不同频率的结果。
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        # 对于偶数下标的维度使用正弦函数来填充P
        self.P[:, :, 0::2] = torch.sin(X)
        # 奇数下标使用余弦函数填充
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        # 将输入张量X与相应长度的位置编码相加。
        # .as_in_ctx(X.ctx) 确保位置编码P与X处于相同的计算上下文（例如GPU或CPU）。
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        
        # 应用dropout层，并返回结果。
        return self.dropout(X)

编码器

#@save
class Encoder(nn.Module):
    """编码器-解码器架构的基本编码器接口"""
    def __init__(self, **kwargs):
        super(Encoder, self).__init__(**kwargs)

    def forward(self, X, *args):
        raise NotImplementedError

# Transformer编码器Block
class EncoderBlock(nn.Module):
    def __init__(self, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, dropout, use_bias=False, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        self.attention = MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout, use_bias)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addnorm = AddNorm(norm_shape, dropout)

    def forward(self, X, valid_lens):
        Y = self.addnorm(X, self.attention(X, X, X, valid_lens))
        return self.addnorm(Y, self.ffn(Y))

# 组装编码器
class TransformerEncoder(Encoder):
    def __init__(self, vocab_size, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, num_layers, dropout, use_bias=False, **kwargs):
        super(TransformerEncoder, self).__init__()
        self.num_hiddens = num_hiddens
        # 堆叠num_layers个EncoderBlock类实例
        # 使用范围-1~1的固定位置编码
        # 因此通过学习得到的输入的嵌入表示的值需要先乘以嵌入维度的平方根进行重新缩放，然后再与位置编码相加。
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        # 添加指定层数的Encoder Block
        for i in range(num_layers):
            self.blks.add_module("block"+str(i),
                EncoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, use_bias))
    
    def forward(self, X, valid_lens, *args):
        # 因为位置编码值在-1和1之间，
        # 因此嵌入值乘以嵌入维度的平方根进行缩放，
        # 然后再与位置编码相加。
        '''
        保持梯度稳定：通过乘以sqrt(num_hiddens)，可以确保输入到后续层的数据具有合理的方差。
        如果直接使用未经缩放的嵌入向量，那么在网络的深层可能会导致梯度消失或爆炸的问题。

        与位置编码匹配：未经缩放的词嵌入可能有较大的数值范围。通过乘以sqrt(num_hiddens)，可以调整词嵌入的大小，
        使其与位置编码的幅度相匹配，从而保证两者相加时不会因为数值差异过大而导致信息失衡。

        防止在多头注意力机制中点积运算后出现过大的数值，这可能导致softmax函数输出过于极端的概率分布，即一个或几个概率接近1，其余接近0。
        通过缩放，可以使softmax函数的输入更均匀地分布在合理范围内，从而获得更平滑的概率分布。

        这个操作的目的并不是为了拉大数值，而是为了平衡词嵌入和位置编码之间的幅度，并帮助模型更好地学习。
        '''
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        # 用于存放注意力权重的数组
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens)
            self.attention_weights[
                i] = blk.attention.attention.attention_weights
        return X

解码器

#@save
class Decoder(nn.Module):
    """编码器-解码器架构的基本解码器接口"""
    def __init__(self, **kwargs):
        super(Decoder, self).__init__(**kwargs)

    def init_state(self, enc_outputs, *args):
        raise NotImplementedError

    def forward(self, X, state):
        raise NotImplementedError

class DecoderBlock(nn.Module):
    # 掩蔽多头解码器自注意力层（第一个子层）中，查询、键和值都来自上一个解码器层的输出。
    # 在训练阶段，其输出序列的所有位置（时间步）的词元都是已知的。
    # 在预测阶段，其输出序列的词元是逐个生成的。
    # 在任何解码器时间步中，只有生成的词元才能用于解码器的自注意力计算中。
    def __init__(self, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, dropout, i, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.i = i
        self.attention1 = MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.attention2 = MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm2 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addnorm3 = AddNorm(norm_shape, dropout)

    def forward(self, X, state):
        # 来自解码器的输出
        enc_outputs, enc_valid_lens = state[0], state[1]
        # state[2][self.i]保存了直到当前时间步为止第i个解码器块解码的所有输出表示。
        '''
        如果state[2][self.i]不为空（即不是None），这意味着这不是该解码器块的第一个时间步，而是后续的一个时间步。在这种情况下，
        需要将之前所有已解码的输出（state[2][self.i]）和新的输入X连接起来，形成一个更长的序列，
        这样自注意力机制就可以考虑到整个已经生成的序列信息来进行下一步的解码。
        如果 state[2][self.i] 是 None，这意味着这是该解码器块第一次被调用，或者在训练过程中每个时间步的状态没有被保留。
        在这种情况下，模型还没有之前的输出可以用来与当前的输入 X 拼接。
        '''
        if state[2][self.i] is None:
            key_values = X
        else:
            '''
            state[2][self.i] 的形状是 (batch_size, prev_seq_len, feature_size) 而 X 的形状是 (batch_size, new_seq_len, feature_size)，
            那么拼接后的 key_values 将会有形状 (batch_size, prev_seq_len + new_seq_len, feature_size)
            '''
            key_values = torch.cat((state[2][self.i], X), axis=1)
        # 保存这个时间步的信息
        state[2][self.i] = key_values
        if self.training:
            batch_size, num_steps, _ = X.shape
            # dec_valid_lens的开头:(batch_size,num_steps)
            # 其中每一行是[1,2,...,num_steps]
            dec_valid_lens = torch.arange(
                1, num_steps + 1, device=X.device).repeat(batch_size, 1)
        else:
            dec_valid_lens = None
        # 掩码自注意力
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
        Y = self.addnorm1(X, X2)
        # 交叉注意力
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), state

class TransformerDecoder(d2l.AttentionDecoder):
    def __init__(self, vocab_size, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, num_layers, dropout, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
           self.blks.add_module("block"+str(i), DecoderBlock(key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, dropout, i))
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, enc_valid_lens, *args):
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]

    def forward(self, X, state):
        # 位置编码
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        # 存放注意力权重
        self._attention_weights = [[None] * len(self.blks) for _ in range (2)]
        for i, blk in enumerate(self.blks):
            X, state = blk(X, state)
            # 存储解码器掩码自注意力权重
            self._attention_weights[0][
                i] = blk.attention1.attention.attention_weights
            # 存储“编码器－解码器”交叉自注意力权重
            self._attention_weights[1][
                i] = blk.attention2.attention.attention_weights
        # 经过全连接层随后输出
        return self.dense(X), state

    @property
    def attention_weights(self):
        return self._attention_weights

进行训练及相关辅助函数

#@save
def sequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项"""
    maxlen = X.size(1)
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value
    return X

#@save
class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""
    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    def forward(self, pred, label, valid_len):
        weights = torch.ones_like(label)
        weights = sequence_mask(weights, valid_len)
        self.reduction='none'
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(
            pred.permute(0, 2, 1), label)
        weighted_loss = (unweighted_loss * weights).mean(dim=1)
        return weighted_loss

def grad_clipping(net, theta):
    """裁剪梯度"""
    if isinstance(net, nn.Module):
        params = [p for p in net.parameters() if p.requires_grad]
    else:
        params = net.params  # 假设net有一个params属性

    # 只计算非空梯度的平方和
    norm = torch.sqrt(sum(torch.sum((p.grad ** 2)) for p in params if p.grad is not None))

    if norm > theta:
        for param in params:
            if param.grad is not None:  # 只对非空梯度进行裁剪
                param.grad[:] *= theta / norm

# 定义超参数
num_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.1, 64, 10
lr, num_epochs, device = 0.005, 200, d2l.try_gpu()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4
key_size, query_size, value_size = 32, 32, 32
norm_shape = [32]
# 读取数据集
train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)

#@save
def train_seq2seq(net, data_iter, lr, num_epochs, tgt_vocab, device):
    """训练序列到序列模型"""
    def xavier_init_weights(m):
        if type(m) == nn.Linear:
            nn.init.xavier_uniform_(m.weight)
        if type(m) == nn.GRU:
            for param in m._flat_weights_names:
                if "weight" in param:
                    nn.init.xavier_uniform_(m._parameters[param])

    net.apply(xavier_init_weights)
    net.to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    loss = MaskedSoftmaxCELoss()
    net.train()
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                     xlim=[10, num_epochs])
    for epoch in range(num_epochs):
        timer = d2l.Timer()
        metric = d2l.Accumulator(2)  # 训练损失总和，词元数量
        for batch in data_iter:
            optimizer.zero_grad()
            X, X_valid_len, Y, Y_valid_len = [x.to(device) for x in batch]
            bos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0],
                          device=device).reshape(-1, 1)
            dec_input = torch.cat([bos, Y[:, :-1]], 1)  # 强制教学
            Y_hat, _ = net(X, dec_input, X_valid_len)
            l = loss(Y_hat, Y, Y_valid_len)
            l.sum().backward()      # 损失函数的标量进行“反向传播”
            grad_clipping(net, 1)
            num_tokens = Y_valid_len.sum()
            optimizer.step()
            with torch.no_grad():
                metric.add(l.sum(), num_tokens)
        if (epoch + 1) % 10 == 0:
            animator.add(epoch + 1, (metric[0] / metric[1],))
    print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} '
        f'tokens/sec on {str(device)}')

encoder = TransformerEncoder(
    len(src_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
decoder = TransformerDecoder(
    len(tgt_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
net = EncoderDecoder(encoder, decoder)
train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

预测

engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {d2l.bleu(translation, fra, k=2):.3f}')