Featured image of post 实例学PyTorch（7）：语言模型（二）——使用Transformer实现词级语言模型
机器学习

实例学PyTorch（7）：语言模型（二）——使用Transformer实现词级语言模型

实例学习PyTorch，使用Transformer实现词级语言模型

背景

这是“实例学PyTorch”系列的第7篇文章。在第6篇文章“实例学PyTorch（6）：语言模型（一）——使用LSTM实现词级语言模型”中，我们简单介绍了使用LSTM实现一个词级语言模型。

LSTM以及其他基于循环神经网络（RNN）的模型在自然语言处理中有着广泛的应用，但是这些模型在处理长距离依赖问题时存在一些问题，例如梯度消失和梯度爆炸。为了解决这些问题，研究者提出了Transformer模型，Transformer模型使用了注意力机制，能够更好地处理长距离依赖问题。本文我们就简单介绍一下Transformer模型，并使用Transformer实现一个简单的词级语言模型。本文参考了PyTorch官方的示例代码中的word_language_model示例。

本文的代码可以在我的GitHub仓库https://github.com/jin-li/pytorch-tutorial中的T06_word_lstm文件夹中找到。

Transformer模型和注意力机制

2017年，Google的研究人员发布了一篇题为《Attention is All You Need》的论文，提出了Transformer模型。由于其在语言模型上的表现十分优秀，它很快取代了LSTM和GRU等循环神经网络模型，成为自然语言处理领域的主流模型。

Transformer模型的结构如下图所示：

Transformer模型结构

可以看到，Transformer模型由编码器（Encoder）和解码器（Decoder）组成。编码器和解码器都是由多个相同的层堆叠而成，每个层包含一个多头自注意力机制（Multi-Head Self-Attention）和一个前馈神经网络（Feed-Forward Neural Network）。

自注意力机制

所谓自注意力机制，就是模型可以同时关注输入序列中的不同位置，从而更好地捕捉输入序列中的信息。多头自注意力机制是指模型可以同时关注输入序列中的不同位置，并且可以通过多个头（head）来学习不同的注意力权重，从而更好地捕捉输入序列中的信息。

举个例子，假设我们有一个输入序列[I, love, you]，我们希望模型能够根据I和you的关系来预测love。在LSTM模型中，模型会逐个处理输入序列中的每个词，但是在Transformer模型中，模型可以同时关注I和you，从而更好地捕捉它们之间的关系。由于Transformer模型可以同时关注输入序列中的不同位置，因此它可以更好地处理长距离依赖问题。但是同时，由于Transformer模型没有循环结构，因此它不能像LSTM那样处理序列中的顺序信息，这就需要在输入序列中加入位置编码（Positional Encoding）来表示词的位置信息。

自注意力机制的计算过程如下：

  1. 首先，我们需要计算查询（Query）、键（Key）和值（Value）的向量表示。这里我们使用输入序列的词向量作为查询、键和值的向量表示。这三个向量表示可以通过一个线性变换得到，即$Q = XW^Q$、$K = XW^K$和$V = XW^V$，其中$X$是输入序列的词向量，$W^Q$、$W^K$和$W^V$是线性变换的权重。

  2. 然后，我们计算注意力分数（Attention Scores）$A$，注意力分数是查询向量$Q$和键向量$K$的点积，再除以$\sqrt{d_k}$，其中$d_k$是查询向量$Q$的维度。即$A = \frac{QK^T}{\sqrt{d_k}}$。

  3. 接着，我们计算注意力权重（Attention Weights）$W$，注意力权重是注意力分数$A$经过Softmax函数得到的。即$W = \text{Softmax}(A)$。

  4. 最后，我们计算自注意力输出（Self-Attention Output）$O$，自注意力输出是注意力权重$W$和值向量$V$的加权和。即$O = W \cdot V$。

在实际应用中，我们通常会使用多头自注意力机制，即将输入序列的词向量分别通过多个线性变换得到多组查询、键和值的向量表示，然后分别计算多组注意力分数、注意力权重和自注意力输出，最后将多组自注意力输出拼接起来，再通过一个线性变换得到最终的输出。

前馈神经网络

前馈神经网络是Transformer模型中的另一个重要组件，它由两个全连接层和一个激活函数组成。前馈神经网络的计算过程如下：

  1. 首先，我们将自注意力输出$O$通过一个全连接层得到中间表示$M$，即$M = O \cdot W_1 + b_1$，其中$W_1$和$b_1$是全连接层的权重和偏置。

  2. 然后，我们将中间表示$M$通过一个激活函数（通常是ReLU）得到前馈神经网络的输出$F$，即$F = \text{ReLU}(M)$。

  3. 最后，我们将前馈神经网络的输出$F$通过另一个全连接层得到最终的输出$O’$，即$O’ = F \cdot W_2 + b_2$，其中$W_2$和$b_2$是全连接层的权重和偏置。

前馈神经网络的作用是对自注意力输出$O$进行非线性变换，从而更好地捕捉输入序列中的信息。

编码器和解码器

自注意力加上前馈神经网络构成了Transformer模型中的一个层，整个Transformer模型是由多个这样的层堆叠而成的，如下图所示：

Transformer层

Transformer模型一般可分为编码器（Encoder）和解码器（Decoder）两部分。编码器用于将输入序列编码成一个上下文向量，解码器用于根据上下文向量生成输出序列。编码器和解码器都是由多个相同的层堆叠而成，每个层包含一个多头自注意力机制和一个前馈神经网络。

编码器的输入是一个词序列，输出是一个上下文向量，解码器的输入是一个上下文向量和一个词序列，输出是一个词序列。在机器翻译等任务中，我们可以将源语言的词序列作为编码器的输入，将目标语言的词序列作为解码器的输入，从而实现源语言到目标语言的翻译。

编码器和解码器的区别在于，解码器在计算自注意力时还会计算编码器的输出的注意力，这是为了更好地捕捉输入序列和输出序列之间的关系。另外，解码器使用的自注意力机制是掩码自注意力机制（Masked Self-Attention），即在计算注意力权重时，解码器只能关注当前位置之前的位置，不能关注当前位置之后的位置。这当然是因为我们的目标就是预测当前和之后的词，自然不能使用之后词的信息，否则就是作弊了。

Transformer模型的分类

虽然标准的Transformer模型包含编码器和解码器，但是在实际应用中，我们也可以只使用编码器或解码器，或者同时使用编码器和解码器。而且同时使用编码器和解码器的模型也并不一定就比只使用编码器或解码器的模型效果更好，这取决于具体的任务和数据集。

编码器模型

编码器模型只包含编码器，用于将输入序列编码成一个上下文向量。编码器模型常用于文本分类、情感分析等任务。比较常用的编码器模型有BERT、RoBERTa等。

解码器模型

解码器模型仅包含解码器，用于根据上下文向量生成输出序列。解码器模型常用于机器翻译、文本生成等任务。比较常用的解码器模型有T5、GPT等。

编码器-解码器模型

编码器-解码器模型同时包含编码器和解码器，用于将输入序列编码成一个上下文向量，并根据上下文向量生成输出序列。编码器-解码器模型通常也用于机器翻译、文本生成等任务。比较常用的编码器-解码器模型有Transformer、BART等。

使用PyTorch中的Transformer实现词级语言模型

PyTorch里面提供了torch.nn.Transformer模块，可以方便地实现Transformer模型。下面我们就将上一篇文章中的LSTM模型替换为Transformer模型，实现一个简单的词级语言模型。

准备数据

和上一篇文章中一样，这里我们依然使用WikiText-2数据集，数据集的下载和处理方法可以参考上一篇文章, 这里不再赘述。

定义模型

PyTorch已经内置了一个Transformer模块torch.nn.Transformer，我们可以直接使用这个模块来实现Transformer模型。但是在使用torch.nn.Transformer模块之前，我们需要先定义一个词嵌入层（Embedding Layer）和一个位置编码层（Positional Encoding Layer）。

词嵌入层

词嵌入的知识在上一篇文章中已经介绍过了，这里我们直接可以使用PyTorch内置的torch.nn.Embedding模块来定义一个词嵌入层。

1

input_embedding = nn.Embedding(vocab_size, embed_size)

其中vocab_size是数据集中字典的大小，embed_size是词嵌入的维度。

位置编码层

所谓位置编码，就是为输入序列中的每个词添加一个位置信息，以便模型更好地捕捉输入序列中的信息。我们可以使用下面的公式来计算位置编码：

$p_{(pos, 2i)} = \sin(pos / 10000^{2i / d_{model}})$

$p_{(pos, 2i+1)} = \cos(pos / 10000^{2i / d_{model}})$

其中$d_{model}$是词嵌入的维度，$pos$是词的位置，$i$是词嵌入的维度的索引。简单来说，上面的位置编码公式的效果就是为每个词的每个维度添加一个正弦或余弦函数的位置信息。使用正余弦函数进行位置编码的原因是，这样可以保证不同位置的词的位置编码之间的距离是相等的，从而更好地捕捉输入序列中的信息。

我们可以使用代码来实现上面的位置编码，这里我们参考了PyTorch官方的示例代码中的PositionalEncoding类：

 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
37
38
39
40
41

class PositionalEncoding(nn.Module):
    r"""Inject some information about the relative or absolute position of the tokens in the sequence.
        The positional encodings have the same dimension as the embeddings, so that the two can be summed.
        Here, we use sine and cosine functions of different frequencies.
    .. math:
        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
        \text{where pos is the word position and i is the embed idx)
    Args:
        d_model: the embed dim (required).
        dropout: the dropout value (default=0.1).
        max_len: the max. length of the incoming sequence (default=5000).
    Examples:
        >>> pos_encoder = PositionalEncoding(d_model)
    """

    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        r"""Inputs of forward function
        Args:
            x: the sequence fed to the positional encoder model (required).
        Shape:
            x: [sequence length, batch size, embed dim]
            output: [sequence length, batch size, embed dim]
        Examples:
            >>> output = pos_encoder(x)
        """

        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

在上面的代码中，我们定义了一个PositionalEncoding类，这个类继承自nn.Module类，用于实现位置编码。在__init__方法中，我们首先定义了一个位置编码矩阵pe，然后计算了正余弦函数的位置编码，最后将位置编码矩阵pe添加到模型的缓冲区中。在forward方法中，我们将输入序列x和位置编码矩阵pe相加，然后通过Dropout层得到最终的输出。

Transformer模型

有了词嵌入层和位置编码层之后，我们就可以使用PyTorch内置的torch.nn.Transformer模块来定义一个Transformer模型了。除了继承nn.Module类，并实现__init__和forward方法，我们还需要定义一个generate_square_subsequent_mask方法，用于生成一个掩码矩阵，这个掩码矩阵在计算注意力权重时会用到。另外，我们再定义一个init_weights方法，用于初始化模型的权重。

 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
37
38

class TransformerModel(nn.Transformer):
    """Container module with an encoder, a recurrent or transformer module, and a decoder."""

    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
        super(TransformerModel, self).__init__(d_model=ninp, nhead=nhead, dim_feedforward=nhid, num_encoder_layers=nlayers)
        self.model_type = 'Transformer'
        self.src_mask = None
        self.pos_encoder = PositionalEncoding(ninp, dropout)

        self.input_emb = nn.Embedding(ntoken, ninp)
        self.ninp = ninp
        self.decoder = nn.Linear(ninp, ntoken)

        self.init_weights()

    def _generate_square_subsequent_mask(self, sz):
        return torch.log(torch.tril(torch.ones(sz,sz)))

    def init_weights(self):
        initrange = 0.1
        nn.init.uniform_(self.input_emb.weight, -initrange, initrange)
        nn.init.zeros_(self.decoder.bias)
        nn.init.uniform_(self.decoder.weight, -initrange, initrange)

    def forward(self, src, has_mask=True):
        if has_mask:
            device = src.device
            if self.src_mask is None or self.src_mask.size(0) != len(src):
                mask = self._generate_square_subsequent_mask(len(src)).to(device)
                self.src_mask = mask
        else:
            self.src_mask = None

        src = self.input_emb(src) * math.sqrt(self.ninp)
        src = self.pos_encoder(src)
        output = self.encoder(src, mask=self.src_mask)
        output = self.decoder(output)
        return F.log_softmax(output, dim=-1)

在上面的代码中，我们定义了一个TransformerModel类，这个类继承自nn.Transformer类，用于实现Transformer模型。在__init__方法中，我们首先调用super(TransformerModel, self).__init__方法初始化父类，然后定义了一个模型类型model_type，一个源掩码src_mask，一个位置编码层pos_encoder，一个词嵌入层input_emb，一个线性层decoder，最后调用init_weights方法初始化模型的权重。

运行模型

有了上述模型之后，我们就可以定义一个训练函数和一个测试函数，然后封装一个加载数据、训练和测试的主函数，来运行我们的模型了。注意，相比上一篇文章中的LSTM模型，Transformer模型需要传入一个注意力头数量参数nhead。我们可以通过调用下面的命令来运行我们的模型：

1

python language_transformer.py --plot

每次运行训练模型的代码时，一旦运行超过一个epoch，原来的模型文件model.pt会被覆盖，所以如果想要保存之前的模型，需要手动将model.pt文件重命名，或者在再次训练时通过--save参数指定保存的模型文件名（可以包含路径）。在一个epoch之后，模型会在验证集上计算损失值，如果验证集上的损失值比之前的最小损失值小，模型会保存在model.pt文件中。

在我的个人电脑上，如果使用GPU（Nvidia GeForce RTX 4060 Ti）训练，每个epoch大约需要28秒，显存占用约548MB；如果使用CPU（Intel i5 9600K）训练，每个epoch大约需要516秒。这里我一共训练了50个epoch，训练集和验证集的损失值如下图所示：

Transformer模型训练和验证损失

可以看到，这个结果和上一篇文章中的LSTM模型的结果差不多，Transformer模型的收敛速度比LSTM模型稍微快一点儿，但最终在验证集上的损失值稍大一点儿。

生成文本

和上一篇文章中一样，我们在训练完模型后，代码会默认把模型保存在model.pt文件中，我们可以加载这个模型，然后使用这个模型生成一些文本。生成文本的方法和上一篇文章中的类似，代码在generate_text.py文件中。我们可以调用下面的命令来生成一些文本：

1

python generate_text.py

一个诡异的情况是，在我的电脑上生成文本只能使用GPU。如果在运行上面的命令时指定了--no-cuda参数使用CPU来，我的电脑会直接崩溃重启，甚至连日志都没来得及输出（至少我没找到）。我能想到的可能原因是，我的电脑CUDA版本是12.5，而PyTorch的CUDA版本是12.1，崩溃重启的原因也可能和这个有关（但我也不确定，因为训练模型时可以不用CUDA，只用CPU，理论上只用CPU生成文本也不会用到CUDA，所以崩溃不应该和CUDA有关）。另外，我的电脑内存有64GB，比GPU显存大得多，所以内存不足也应该不是问题。如果有人遇到过类似的问题，欢迎留言告诉我，我们可以讨论一下，谢谢！

Transformer模型生成文本的速度比LSTM模型略慢，在我的电脑上用GPU生成1000个词大约需要6秒。由于使用CPU生成文本会导致电脑崩溃重启，因此我没有测试CPU生成文本的速度。

总结

本文中，我们简单介绍了Transformer模型和注意力机制，然后参照PyTorch官方的“word_language_model”示例，使用PyTorch中的torch.nn.Transformer模块实现了一个简单的词级语言模型。我们还介绍了Transformer模型的分类，包括编码器模型、解码器模型和编码器-解码器模型，模型的训练结果和LSTM差不多。最后，我们可以使用训练好的Transformer模型生成一些文本。

comments powered by Disqus