Featured image of post 实例学PyTorch(6):语言模型——使用LSTM实现词级语言模型(一)

实例学PyTorch(6):语言模型——使用LSTM实现词级语言模型(一)

实例学习PyTorch,使用LSTM实现词级语言模型

背景

这是“实例学PyTorch”系列的第6篇文章。在第4篇文章和第5篇文章中,我们介绍了序列预测问题,并用RNN、GRU和LSTM实现了对正弦函数的预测。

我们在第4篇文章中提到过,除了Sine函数这种时间序列,序列数据还可以是语言模型中的单词序列,本文就来简单介绍使用LSTM实现一个词级语言模型。

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

语言模型

语言模型是自然语言处理中的一个重要问题,它是用来评估一个句子的概率的模型。语言模型可以用来预测下一个单词是什么,也可以用来生成一个句子。语言模型在机器翻译、语音识别、文本生成等任务中都有广泛的应用。

语言模型也是一个序列预测问题,即我们在预测下一个单词时,不仅需要考虑当前单词,还需要考虑前面的单词。例如,当前的单词是“苹果”,我们在预测下一个单词时,首先需要考虑这个“苹果”是指水果还是公司,这就需要用到前面的单词提供的信息。例如,前面的单词里有“吃”、“香蕉”、“梨”之类的,那么“苹果”很可能是指水果;如果前面的单词有“手机”、“电脑”、“乔布斯”之类的,那么“苹果”很可能是指公司。

在前面的文章中,我们简单介绍过,解决序列预测的问题可以使用循环神经网络(RNN)及其变种,例如长短期记忆网络(LSTM)和门控循环单元(GRU)。在PyTorch官方的示例代码中有一个使用RNN等和Transformer实现语言模型的示例,这里我们将基于这个示例代码,使用LSTM实现一个简单的词级语言模型。另外,本文中的部分图片和代码参考和引用了YouTube博主Donato Capitella的视频《LLM Chronicles #4.4: Building a Word-Level Language Model in PyTorch using RNNs》和所附代码中的内容。

当然,近年来,随着Transformer的出现,语言模型的效果得到了极大的提升,但是LSTM作为一种经典的循环神经网络,仍然有着广泛的应用。所以这里我们先介绍用LSTM来实现一个简单的语言模型。在后续的文章中,我们将介绍Transformer及其变种,以及如何使用Transformer实现语言模型。

语言模型分为字符级、词级、次词级等,它们的区别在于对序列的划分不同。字符级即以单个字符为单位,词级以单词为单位,次词级以词素为单位。所谓次词级,是指将一个单词按照词素划分,例如“joyfulness”一词可以划分为“joy”、“ful”和“ness“三个词素。

不同层级的分词示例

词级语言模型

词级语言模型是一种用来预测下一个单词是什么的模型。词级语言模型的输入是一个单词序列,输出也是一个单词序列。在训练阶段,我们将一个句子中的前面的单词作为输入,后面的单词作为输出,通过最小化预测单词和真实单词之间的差异来训练模型。在测试阶段,我们可以使用模型来预测下一个单词。

用LSTM实现词级语言模型的流程和用LSTM或RNN实现正弦函数预测的流程类似:

  1. 准备数据:我们需要将单词序列转换为整数序列,以便输入模型。
  2. 构建模型:我们需要构建一个LSTM模型,用于预测下一个单词。
  3. 训练模型:我们需要使用数据集训练模型,使模型能够预测下一个单词。
  4. 测试模型:我们需要使用模型来预测下一个单词。

准备数据

这里我们使用的数据库是WikiText-2,它是一个常用的语言模型数据集,包含了一些维基百科的词条,其内容可以在这里查看:WikiText-2

WikiText-2

语言模型和正弦函数序列模型不同之处在于,正弦函数序列的输入是数字,而语言模型的输入是单词。PyTorch能处理的都是数学张量,因此我们需要将单词转换为数字。

数据预处理

WikiText-2数据集是一个文本文件,我们需要将文本文件转换为单词序列。PyTorch本来有一个torchtext库,可以用来处理文本数据,但在2024年4月之后,这个库就停止维护了。所以这里我们需要自己写一些代码来预处理训练了数据。

预处理的基本思路是:

  1. 读取全部文本,用一个字典来保存所有不重复的单词。字典的键是单词,值是单词对应的编号。
  2. 将原来文本中的单词全部替换为编号。

这样我们就将训练的文本数据转换为了一个整数序列。这里我们用的代码来自PyTorch官方的示例代码:

 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
42
43
44
45
46
47
48
import os
from io import open
import torch

class Dictionary(object):
    def __init__(self):
        self.word2idx = {}
        self.idx2word = []

    def add_word(self, word):
        if word not in self.word2idx:
            self.idx2word.append(word)
            self.word2idx[word] = len(self.idx2word) - 1
        return self.word2idx[word]

    def __len__(self):
        return len(self.idx2word)


class Corpus(object):
    def __init__(self, path):
        self.dictionary = Dictionary()
        self.train = self.tokenize(os.path.join(path, 'train.txt'))
        self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
        self.test = self.tokenize(os.path.join(path, 'test.txt'))

    def tokenize(self, path):
        """Tokenizes a text file."""
        assert os.path.exists(path)
        # Add words to the dictionary
        with open(path, 'r', encoding="utf8") as f:
            for line in f:
                words = line.split() + ['<eos>']
                for word in words:
                    self.dictionary.add_word(word)

        # Tokenize file content
        with open(path, 'r', encoding="utf8") as f:
            idss = []
            for line in f:
                words = line.split() + ['<eos>']
                ids = []
                for word in words:
                    ids.append(self.dictionary.word2idx[word])
                idss.append(torch.tensor(ids).type(torch.int64))
            ids = torch.cat(idss)

        return ids

这里我们定义了一个Dictionary类和一个Corpus类,Dictionary类用来保存单词和编号的对应关系,Corpus类用来读取文本文件,并将文本文件转换为整数序列。这里我们分别转换了训练集、验证集和测试集。

经过处理后,我们得到了训练集、验证集和测试集的3个整数序列,其长度分别为2088628、217646和245569。字典的大小为33278。

数据分批

经过上一步的处理,我们得到了一个非常长的整数序列,为了便于训练,我们将这个序列划分为若干个小的序列。例如对于一个长度为10000的序列,我们想分为20个批次,那么每个批次的长度就是500。这样我们就得到了一个形状为$500 \times 20$的矩阵。

注意:

  • 转换后的矩阵是列优先的,即第一列是第一个批次的数据,第二列是第二个批次的数据,以此类推。
  • 若数据长度不能整除批次长度,我们将舍弃多余的数据。
  • 各个批次之间是相互独立的,即第一个批次的最后一个数据和第二个批次的第一个数据之间没有关系。这也意味着分批之后的数据会丢失一些上下文信息。

分批化的代码很简单:

1
2
3
4
5
6
7
8
def batchify(data, bsz):
    # Work out how cleanly we can divide the dataset into bsz parts.
    nbatch = data.size(0) // bsz
    # Trim off any extra elements that wouldn't cleanly fit (remainders).
    data = data.narrow(0, 0, nbatch * bsz)
    # Evenly divide the data across the bsz batches.
    data = data.view(bsz, -1).t().contiguous()
    return data

词嵌入

我们在之前的文章中,我们使用过one-hot对标签进行编码。但是对于单词,one-hot编码就不太合适了,因为单词的数量太多(例如在Wikitext-2的字典中有33278个单词),one-hot编码会导致维度过高,计算量过大。因此我们需要使用另一种方法来表示单词,将单词映射到一个低维空间。这个步骤就叫做词嵌入(word embedding)。

简单来说,词嵌入的想法是,使用若干个特征来表示一个单词,这些特征可以是单词的词性、情感、语义等信息。例如,我们选取“活物(living being)“、“猫科(feline)”、“人类(human)”、“性别(gender)”、“皇家(royalty)”、“动词(verb)”、“复数(plural)”7个特征,对于一个单词,我们分别用一个-1到1之间的数值来描述这些特征的程度,把表示所有特征的值组合起来得到一个7维向量,这个向量就是这个单词的词嵌入。例如对于单词“man”,我们可以向量$[0.6, -0.2, 0.8, 0.9, -0.1, -0.9, -0.7]$来表示单词“man”的词嵌入。

类似地,我们可以将字典里的单词都表示为词嵌入。下图是几个例子:

词嵌入示例

词嵌入是语言模型中一个非常重要的概念,因为它可以帮助我们的语言模型更好地表示单词的语义信息,也可以清楚地看到单词之间的关系。例如,我们将“king”和“man”相减,再加上“woman”,得到的结果应该和“queen”很接近。

当然,实际应用中,我们不会仅使用7个特征,而是一般会使用几百个特征,这样可以全面地表示单词的语义信息。例如对于本问题,我们选择使用200个特征,那么我们需要将字典中的33278个单词映射到一个200维的空间。如果用矩阵来表示,这个词嵌入矩阵就是一个$33278 \times 200$的矩阵。

这个想法看上去挺合理的,但是如何得到这个词嵌入矩阵呢?我们也可以通过训练一个神经网络来得到这个词嵌入矩阵。这个神经网络的输入是一个单词的编号,输出是这个单词的词嵌入。这个神经网络的结构可以是一个全连接层,也可以是一个卷积层,也可以是一个LSTM层。这个神经网络的训练目标是最小化预测的词嵌入和真实的词嵌入之间的差异。这个神经网络的训练过程和语言模型的训练过程是类似的,只是输入和输出不同。

当然,这里我们并不需要自己来训练这个词嵌入矩阵,因为研究者训练过很多这样的词嵌入矩阵,我们可以直接使用这些词嵌入矩阵。这些词嵌入矩阵可以是通用的,也可以是针对某个特定任务训练的。在PyTorch中,我们可以使用torch.nn.Embedding来加载这些词嵌入矩阵。

1
torch.nn.Embedding(ntoken, emsize)

其中第一个参数是字典的大小,第二个参数是词嵌入的维度。这个函数会返回一个Embedding对象,我们将这个对象作为LSTM模型中的一层。

LSTM模型

有了词嵌入之后,我们就可以构建LSTM模型了。这里我们以PyTorch中的torch.nn.LSTM为基础,构建一个简单的LSTM模型。这个模型的结构如下:

 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
import torch.nn as nn
import torch.nn.functional as F

class LanguageLSTM(nn.Module):
    """Container module with an encoder, a recurrent module, and a decoder."""

    def __init__(self, rnn_type, ntoken, ninp, nhid, nlayers, dropout=0.5, tie_weights=False):
        super(LanguageLSTM, self).__init__()
        self.ntoken = ntoken
        self.drop = nn.Dropout(dropout)
        self.encoder = nn.Embedding(ntoken, ninp)
        self.rnn = getattr(nn, rnn_type)(ninp, nhid, nlayers, dropout=dropout)
        self.decoder = nn.Linear(nhid, ntoken)

        self.init_weights()

        self.rnn_type = rnn_type
        self.nhid = nhid
        self.nlayers = nlayers

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

    def forward(self, input, hidden):
        emb = self.drop(self.encoder(input))
        output, hidden = self.rnn(emb, hidden)
        output = self.drop(output)
        decoded = self.decoder(output)
        decoded = decoded.view(-1, self.ntoken)
        return F.log_softmax(decoded, dim=1), hidden

    def init_hidden(self, bsz):
        weight = next(self.parameters())
        return (weight.new_zeros(self.nlayers, bsz, self.nhid),
                weight.new_zeros(self.nlayers, bsz, self.nhid))

这个模型的结构和之前的LSTM模型类似,只是这里我们使用了词嵌入层。这个模型的输入是一个整数序列,输出是一个概率分布,表示下一个单词是哪个单词的概率。这个模型的训练目标是最小化预测的概率分布和真实的概率分布之间的差异。

训练和测试模型

有了数据和模型之后,我们就可以开始训练模型了。训练模型的代码和之前的LSTM模型类似,只是这里的损失函数criterion我们一般使用CrossEntropyLoss

除了在训练时在命令行输出损失值,我们在最后也把一个epoch内所有批次的损失值加权平均后返回,便于之后的绘图。

 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
def train(device, model, epoch, train_data, batch_size, criterion, lr, log_interval, seq_len):
    model.train()
    total_loss = 0.
    loss_all = []
    data_cnt = []
    start_time = time.time()
    hidden = model.init_hidden(batch_size)
    for batch, i in enumerate(range(0, train_data.size(0) - 1, seq_len)):
        data, targets = get_batch(train_data, i)
        data, targets = data.to(device), targets.to(device)
        model.zero_grad()
        hidden = repackage_hidden(hidden)
        output, hidden = model(data, hidden)
        loss = criterion(output, targets)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.25)
        for p in model.parameters():
            p.data.add_(p.grad, alpha=-lr)
        total_loss += loss.item()
        loss_all.append(loss.item())
        data_cnt.append(len(data))
        if batch % log_interval == 0 and batch > 0:
            cur_loss = total_loss / log_interval
            elapsed = time.time() - start_time
            print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:02.2f} | ms/batch {:5.2f} | '
                  'loss {:5.2f} | ppl {:8.2f}'.format(
                epoch, batch, len(train_data) // seq_len, lr,
                elapsed * 1000 / log_interval, cur_loss, math.exp(cur_loss)))
            total_loss = 0
            start_time = time.time()
    return np.average(loss_all, weights=data_cnt)

测试模型的代码也和之前的LSTM模型类似,这里不再赘述。

模型性能

全部的代码参见我的GitHub仓库https://github.com/jin-li/pytorch-tutorial中的T06_word_lstm文件夹中。在配置好环墶后,我们可以运行language_lstm.py来训练模型。

1
python language_lstm.py

在我的个人电脑上,如果使用GPU(Nvidia GeForce RTX 4060 Ti)训练,每个epoch大约需要26秒,显存占用约540MB;如果使用CPU(Intel i5 9600K)训练,每个epoch大约需要506秒。这里我一共训练了50个epoch,训练集和验证集的损失值如下图所示:

训练集和验证集的损失值

可以看到,在训练20个epoch之后,模型的损失值就基本稳定在4.1左右,验证集的损失值也基本稳定在4.7左右。这说明模型的泛化能力大致还行。而且我们的训练数据不算多,所以20个epoch的训练就基本够了。

使用模型生成文本

模型训练完成时,训练好的模型会作为model.pt文件保存在当前目录下。我们可以使用这个模型来生成文本。生成文本的代码如下:

 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
def generate_text(device, checkpoint, data_source, words, temperature, log_interval):
    
    with open(checkpoint, 'rb') as f:
        model = torch.load(f, map_location=device)
    model.eval()

    corpus = data.Corpus(data_source)
    ntokens = len(corpus.dictionary)

    hidden = model.init_hidden(1)
    input = torch.randint(ntokens, (1, 1), dtype=torch.long).to(device)

    generated_text = []
    with torch.no_grad():  # no tracking history
        for i in range(words):
            output, hidden = model(input, hidden)
            word_weights = output.squeeze().div(temperature).exp().cpu()
            word_idx = torch.multinomial(word_weights, 1)[0]
            input.fill_(word_idx)

            word = corpus.dictionary.idx2word[word_idx]
            generated_text.append(word)

            if i % log_interval == 0:
                print('| Generated {}/{} words'.format(i, words))
    
    return generated_text

这个函数里的checkpoint是模型文件的路径,data_source是数据集的路径,words是生成的单词数量,temperature是控制生成文本的多样性的参数,log_interval是每隔多少单词输出一次。调用这个函数生成文本并保存的完整代码参见generate_text.py。运行这个代码:

1
python generate_text.py

我们可以得到生成的文本,这里我生成了1000个单词,其中的一部分如下:

1
– <unk> , a year then with the software . It usually was sold for nearly half the day time . For this reason , the Nevermind run surpassed and a new group of canned <unk> . It had benefited from the unhealthy content , which have been leveled on the <unk> 's gates through the design the effects products associated with other birds and tested stewardship of those articles , ranging from an upright system with <unk> <unk> .

由于我们训练使用的WikiText-2数据集是维基百科的词条,其中包含了很多非ASCII字符,因此生成的文本中可能会有一些<unk>字符,这是因为这些字符不在我们的字典中。另外,这里的文本是随机生成的,因此可能不通顺。我们可以调整temperature参数来控制生成文本的多样性,temperature越大,生成的文本越多样化,temperature越小,生成的文本越保守。

总结

本文介绍了使用LSTM实现一个简单的词级语言模型。语言模型是自然语言处理中的一个重要问题,它可以用来预测下一个单词是什么,也可以用来生成一个句子。语言模型是一个序列预测问题,这里我们使用LSTM来解决这个问题,当然也可以使用其他的RNN模型。

使用RNN及其变种来实现语言模型固然可以得到不错的效果,但其在复杂任务上的表现仍然有很大的不足。近年来,随着Transformer的出现,语言模型的效果得到了极大的提升。在后续的文章中,我们将介绍Transformer及其变种,以及如何使用Transformer实现语言模型。

comments powered by Disqus