CBoW-PyTorch

训练 CBoW 模型#

本笔记本是 AI for Beginners Curriculum 的一部分

在这个例子中，我们将学习如何训练 CBoW 语言模型，以获得我们自己的 Word2Vec 嵌入空间。我们将使用 AG News 数据集作为文本来源。

In [ ]:
import torch
import torchtext
import os
import collections
import builtins
import random
import numpy as np
In [ ]:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

首先，让我们加载数据集并定义分词器和词汇表。我们将 vocab_size 设置为 5000 以稍微限制计算量。

In [ ]:
def load_dataset(ngrams = 1, min_freq = 1, vocab_size = 5000 , lines_cnt = 500):
    tokenizer = torchtext.data.utils.get_tokenizer('basic_english')
    print("Loading dataset...")
    test_dataset, train_dataset  = torchtext.datasets.AG_NEWS(root='./data')
    train_dataset = list(train_dataset)
    test_dataset = list(test_dataset)
    classes = ['World', 'Sports', 'Business', 'Sci/Tech']
    print('Building vocab...')
    counter = collections.Counter()
    for i, (_, line) in enumerate(train_dataset):
        counter.update(torchtext.data.utils.ngrams_iterator(tokenizer(line),ngrams=ngrams))
        if i == lines_cnt:
            break
    vocab = torchtext.vocab.Vocab(collections.Counter(dict(counter.most_common(vocab_size))), min_freq=min_freq)
    return train_dataset, test_dataset, classes, vocab, tokenizer
In [ ]:
train_dataset, test_dataset, _, vocab, tokenizer = load_dataset()
Loading dataset...
Building vocab...
In [ ]:
def encode(x, vocabulary, tokenizer = tokenizer):
    return [vocabulary[s] for s in tokenizer(x)]

CBoW 模型#

CBoW 通过 $2N$ 个邻近词来预测一个单词。例如，当 $N=1$ 时，我们可以从句子 I like to train networks 中得到以下配对：(like,I)、(I, like)、(to, like)、(like,to)、(train,to)、(to, train)、(networks, train)、(train,networks)。这里，第一个单词是作为输入的邻近词，第二个单词是我们要预测的目标词。

为了构建一个预测下一个单词的网络，我们需要将邻近词作为输入，并输出单词编号。CBoW 网络的架构如下：

  • 输入单词会通过嵌入层处理。这个嵌入层就是我们的 Word2Vec 嵌入，因此我们会将其单独定义为 embedder 变量。在这个例子中，我们将使用嵌入维度为 30，尽管你可能希望尝试更高的维度（真实的 Word2Vec 通常使用 300）。
  • 嵌入向量接着会传递到一个线性层，该层用于预测输出单词。因此它包含 vocab_size 个神经元。

对于输出，如果我们使用 CrossEntropyLoss 作为损失函数，那么我们只需要提供单词编号作为期望结果，而不需要进行独热编码。

In [ ]:
vocab_size = len(vocab)

embedder = torch.nn.Embedding(num_embeddings = vocab_size, embedding_dim = 30)
model = torch.nn.Sequential(
    embedder,
    torch.nn.Linear(in_features = 30, out_features = vocab_size),
)

print(model)
Sequential(
  (0): Embedding(5002, 30)
  (1): Linear(in_features=30, out_features=5002, bias=True)
)

准备训练数据#

现在让我们编写一个主函数，用于从文本中计算CBoW词对。这个函数将允许我们指定窗口大小，并返回一组词对——输入词和输出词。请注意，这个函数既可以用于单词，也可以用于向量/张量——这将使我们能够在传递给to_cbow函数之前对文本进行编码。

In [ ]:
def to_cbow(sent,window_size=2):
    res = []
    for i,x in enumerate(sent):
        for j in range(max(0,i-window_size),min(i+window_size+1,len(sent))):
            if i!=j:
                res.append([sent[j],x])
    return res

print(to_cbow(['I','like','to','train','networks']))
print(to_cbow(encode('I like to train networks', vocab)))
[['like', 'I'], ['to', 'I'], ['I', 'like'], ['to', 'like'], ['train', 'like'], ['I', 'to'], ['like', 'to'], ['train', 'to'], ['networks', 'to'], ['like', 'train'], ['to', 'train'], ['networks', 'train'], ['to', 'networks'], ['train', 'networks']]
[[232, 172], [5, 172], [172, 232], [5, 232], [0, 232], [172, 5], [232, 5], [0, 5], [1202, 5], [232, 0], [5, 0], [1202, 0], [5, 1202], [0, 1202]]

让我们准备训练数据集。我们将遍历所有新闻，调用to_cbow获取单词对列表，并将这些对添加到X和Y中。为了节省时间，我们只考虑前1万条新闻——如果你有更多时间等待，并希望获得更好的嵌入，可以轻松去掉这个限制 :)

In [ ]:
X = []
Y = []
for i, x in zip(range(10000), train_dataset):
    for w1, w2 in to_cbow(encode(x[1], vocab), window_size = 5):
        X.append(w1)
        Y.append(w2)

X = torch.tensor(X)
Y = torch.tensor(Y)

我们还将把这些数据转换为一个数据集，并创建数据加载器：

In [ ]:
class SimpleIterableDataset(torch.utils.data.IterableDataset):
    def __init__(self, X, Y):
        super(SimpleIterableDataset).__init__()
        self.data = []
        for i in range(len(X)):
            self.data.append( (Y[i], X[i]) )
        random.shuffle(self.data)

    def __iter__(self):
        return iter(self.data)

我们还将把这些数据转换为一个数据集，并创建数据加载器：

In [ ]:
ds = SimpleIterableDataset(X, Y)
dl = torch.utils.data.DataLoader(ds, batch_size = 256)

现在让我们进行实际训练。我们将使用SGD优化器，并设置较高的学习率。您也可以尝试使用其他优化器，例如Adam。我们将从训练10个周期开始——如果您希望获得更低的损失，可以重新运行此单元格。

In [ ]:
def train_epoch(net, dataloader, lr = 0.01, optimizer = None, loss_fn = torch.nn.CrossEntropyLoss(), epochs = None, report_freq = 1):
    optimizer = optimizer or torch.optim.Adam(net.parameters(), lr = lr)
    loss_fn = loss_fn.to(device)
    net.train()

    for i in range(epochs):
        total_loss, j = 0, 0, 
        for labels, features in dataloader:
            optimizer.zero_grad()
            features, labels = features.to(device), labels.to(device)
            out = net(features)
            loss = loss_fn(out, labels)
            loss.backward()
            optimizer.step()
            total_loss += loss
            j += 1
        if i % report_freq == 0:
            print(f"Epoch: {i+1}: loss={total_loss.item()/j}")

    return total_loss.item()/j
In [ ]:
train_epoch(net = model, dataloader = dl, optimizer = torch.optim.SGD(model.parameters(), lr = 0.1), loss_fn = torch.nn.CrossEntropyLoss(), epochs = 10)
Epoch: 1: loss=5.664632366860172
Epoch: 2: loss=5.632101973960962
Epoch: 3: loss=5.610399051405015
Epoch: 4: loss=5.594621561080262
Epoch: 5: loss=5.582538017415446
Epoch: 6: loss=5.572900234519603
Epoch: 7: loss=5.564951676341915
Epoch: 8: loss=5.558288112064614
Epoch: 9: loss=5.552576955031129
Epoch: 10: loss=5.547634165194347

5.547634165194347

尝试使用 Word2Vec#

为了使用 Word2Vec，让我们提取与词汇表中所有单词对应的向量：

In [ ]:
vectors = torch.stack([embedder(torch.tensor(vocab[s])) for s in vocab.itos], 0)

让我们看看，例如，单词Paris是如何编码成一个向量的：

In [ ]:
paris_vec = embedder(torch.tensor(vocab['paris']))
print(paris_vec)
tensor([-0.0915,  2.1224, -0.0281, -0.6819,  1.1219,  0.6458, -1.3704, -1.3314,
        -1.1437,  0.4496,  0.2301, -0.3515, -0.8485,  1.0481,  0.4386, -0.8949,
         0.5644,  1.0939, -2.5096,  3.2949, -0.2601, -0.8640,  0.1421, -0.0804,
        -0.5083, -1.0560,  0.9753, -0.5949, -1.6046,  0.5774],
       grad_fn=<EmbeddingBackward>)

使用Word2Vec寻找同义词非常有趣。以下函数将返回与给定输入最接近的n个单词。为了找到它们，我们计算$|w_i - v|$的范数，其中$v$是与输入单词对应的向量，$w_i$是词汇表中第$i$个单词的编码。然后我们对数组进行排序，并使用argsort返回对应的索引，取列表的前n个元素，这些元素编码了词汇表中最接近单词的位置。

In [ ]:
def close_words(x, n = 5):
  vec = embedder(torch.tensor(vocab[x]))
  top5 = np.linalg.norm(vectors.detach().numpy() - vec.detach().numpy(), axis = 1).argsort()[:n]
  return [ vocab.itos[x] for x in top5 ]

close_words('microsoft')
['microsoft', 'quoted', 'lp', 'rate', 'top']
In [ ]:
close_words('basketball')
['basketball', 'lot', 'sinai', 'states', 'healthdaynews']
In [ ]:
close_words('funds')
['funds', 'travel', 'sydney', 'japan', 'business']

要点#

通过使用诸如CBoW这样的巧妙技术，我们可以训练Word2Vec模型。你也可以尝试训练skip-gram模型，该模型通过给定中心词来预测邻近词，看看它的表现如何。

免责声明：
本文档使用AI翻译服务 Co-op Translator 进行翻译。尽管我们努力确保翻译的准确性，但请注意，自动翻译可能包含错误或不准确之处。应以原始语言的文档作为权威来源。对于关键信息，建议使用专业人工翻译。我们对因使用此翻译而引起的任何误解或误读不承担责任。