使用 PyTorch 实现 Word2Vec 中Skip-gram 模型

首先创建了一个使用 Word2VecDataset 类自定义的数据集，用于生成训练数据。然后，定义了 Skip-gram 模型，并使用交叉熵损失函数和 Adam 优化器进行训练。

在每个训练周期中，遍历数据加载器，对每个批次进行前向传播、计算损失、反向传播和权重更新。最后，得到训练得到的词向量，并可以使用 word_vector 来获取特定单词的词向量表示。

确保在运行之前安装 PyTorch，可以使用 pip install torch 来安装它。请注意，如果可用的话，代码将在 GPU 上运行。如果没有 GPU，请将 .to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')) 部分删除，并在 CPU 上运行。

以下是使用 PyTorch 实现 Skip-gram 模型的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
 
# Hyperparameters
embedding_dim = 100
window_size = 2
learning_rate = 0.001
epochs = 100
batch_size = 32
 
# Example corpus
corpus = [['I', 'enjoy', 'playing', 'football', 'with', 'my', 'friends'],
          ['We', 'like', 'to', 'play', 'tennis', 'on', 'weekends'],
          ['She', 'is', 'a', 'good', 'dancer']]
 
# Create vocabulary
vocab = list(set([word for sentence in corpus for word in sentence]))
vocab_size = len(vocab)
word2idx = {word: idx for idx, word in enumerate(vocab)}
idx2word = {idx: word for idx, word in enumerate(vocab)}
 
# Generate training data
class Word2VecDataset(Dataset):
    def __init__(self, corpus, word2idx):
        self.data = []
        for sentence in corpus:
            word_indices = [word2idx[word] for word in sentence]
            for center_word_idx, center_word in enumerate(word_indices):
                for context_word_idx in range(max(0, center_word_idx - window_size), min(center_word_idx + window_size + 1, len(word_indices))):
                    if context_word_idx != center_word_idx:
                        context_word = word_indices[context_word_idx]
                        self.data.append((center_word, context_word))
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]
 
dataset = Word2VecDataset(corpus, word2idx)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
# Define Skip-gram model
class SkipGramModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(SkipGramModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.linear = nn.Linear(embedding_dim, vocab_size)
        
    def forward(self, center_word):
        embedded = self.embedding(center_word)
        output = self.linear(embedded)
        return output
 
model = SkipGramModel(vocab_size, embedding_dim)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
 
# Training
for epoch in range(epochs):
    running_loss = 0.0
    for i, (center_word, context_word) in enumerate(dataloader):
        optimizer.zero_grad()
        
        center_word = center_word.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
        context_word = context_word.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
        
        output = model(center_word)
        loss = criterion(output, context_word)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        
    average_loss = running_loss / len(dataloader)
    print(f'Epoch {epoch+1}/{epochs}, Loss: {average_loss:.4f}')
 
# Get trained word embeddings
trained_embeddings = model.embedding.weight.data.numpy()
 
# Example usage - Getting word vector for a word
word = 'football'
word_vector = trained_embeddings[word2idx[word]]
print(f"Word vector for '{word}': {word_vector}")

运行结果如下：

Epoch 1/100, Loss: 3.1324
Epoch 2/100, Loss: 3.0791
Epoch 3/100, Loss: 2.9902
Epoch 4/100, Loss: 2.9392
Epoch 5/100, Loss: 2.8870
Epoch 6/100, Loss: 2.8166
Epoch 7/100, Loss: 2.7615
Epoch 8/100, Loss: 2.7017
Epoch 9/100, Loss: 2.6500
Epoch 10/100, Loss: 2.5993
Epoch 11/100, Loss: 2.5496
Epoch 12/100, Loss: 2.5013
Epoch 13/100, Loss: 2.4621
Epoch 14/100, Loss: 2.4079
Epoch 15/100, Loss: 2.3660
Epoch 16/100, Loss: 2.3229
Epoch 17/100, Loss: 2.2795
Epoch 18/100, Loss: 2.2398
Epoch 19/100, Loss: 2.1998
Epoch 20/100, Loss: 2.1582
Epoch 21/100, Loss: 2.1278
Epoch 22/100, Loss: 2.1023
Epoch 23/100, Loss: 2.0569
Epoch 24/100, Loss: 2.0245
Epoch 25/100, Loss: 1.9936
Epoch 26/100, Loss: 1.9639
Epoch 27/100, Loss: 1.9344
Epoch 28/100, Loss: 1.9137
Epoch 29/100, Loss: 1.8888
Epoch 30/100, Loss: 1.8586
Epoch 31/100, Loss: 1.8352
Epoch 32/100, Loss: 1.8200
Epoch 33/100, Loss: 1.7815
Epoch 34/100, Loss: 1.7685
Epoch 35/100, Loss: 1.7531
Epoch 36/100, Loss: 1.7209
Epoch 37/100, Loss: 1.7049
Epoch 38/100, Loss: 1.6881
Epoch 39/100, Loss: 1.6775
Epoch 40/100, Loss: 1.6517
Epoch 41/100, Loss: 1.6390
Epoch 42/100, Loss: 1.6238
Epoch 43/100, Loss: 1.6077
Epoch 44/100, Loss: 1.5939
Epoch 45/100, Loss: 1.5745
Epoch 46/100, Loss: 1.5703
Epoch 47/100, Loss: 1.5574
Epoch 48/100, Loss: 1.5458
Epoch 49/100, Loss: 1.5308
Epoch 50/100, Loss: 1.5215
Epoch 51/100, Loss: 1.5122
Epoch 52/100, Loss: 1.4988
Epoch 53/100, Loss: 1.4958
Epoch 54/100, Loss: 1.4773
Epoch 55/100, Loss: 1.4746
Epoch 56/100, Loss: 1.4618
Epoch 57/100, Loss: 1.4560
Epoch 58/100, Loss: 1.4506
Epoch 59/100, Loss: 1.4380
Epoch 60/100, Loss: 1.4266
Epoch 61/100, Loss: 1.4257
Epoch 62/100, Loss: 1.4148
Epoch 63/100, Loss: 1.4090
Epoch 64/100, Loss: 1.4070
Epoch 65/100, Loss: 1.3940
Epoch 66/100, Loss: 1.3890
Epoch 67/100, Loss: 1.3846
Epoch 68/100, Loss: 1.3813
Epoch 69/100, Loss: 1.3738
Epoch 70/100, Loss: 1.3717
Epoch 71/100, Loss: 1.3681
Epoch 72/100, Loss: 1.3594
Epoch 73/100, Loss: 1.3593
Epoch 74/100, Loss: 1.3504
Epoch 75/100, Loss: 1.3447
Epoch 76/100, Loss: 1.3439
Epoch 77/100, Loss: 1.3397
Epoch 78/100, Loss: 1.3315
Epoch 79/100, Loss: 1.3260
Epoch 80/100, Loss: 1.3253
Epoch 81/100, Loss: 1.3229
Epoch 82/100, Loss: 1.3215
Epoch 83/100, Loss: 1.3148
Epoch 84/100, Loss: 1.3160
Epoch 85/100, Loss: 1.3072
Epoch 86/100, Loss: 1.3105
Epoch 87/100, Loss: 1.3104
Epoch 88/100, Loss: 1.3018
Epoch 89/100, Loss: 1.2912
Epoch 90/100, Loss: 1.2950
Epoch 91/100, Loss: 1.2938
Epoch 92/100, Loss: 1.2951
Epoch 93/100, Loss: 1.2859
Epoch 94/100, Loss: 1.2902
Epoch 95/100, Loss: 1.2840
Epoch 96/100, Loss: 1.2748
Epoch 97/100, Loss: 1.2840
Epoch 98/100, Loss: 1.2763
Epoch 99/100, Loss: 1.2772
Epoch 100/100, Loss: 1.2746

Word vector for 'football':

[-1.2727762   0.8401019  -0.5115612   2.0667355   1.1854529  -0.7444803
 -1.9658612  -1.0488677   0.98938674 -1.1675086   1.582392    1.7414839
 -0.4892138  -1.2149098   0.15343344 -1.8318586   0.41794038  0.25481498
  0.6008032  -0.23904797  0.80143225 -1.0495795  -1.0174142  -0.01827855
  2.7477944  -0.9574399   1.025569    2.4843202  -0.2796719  -0.4390253
 -1.4423424  -1.8073392   0.1897556   0.90259725  2.7565296  -0.28331178
 -1.8443514   0.77545553 -1.0289538   0.71483964  1.1801128  -0.22635305
  0.5960759   0.6690206  -1.9100318   1.2388043  -0.68522704  0.92120373
  1.0252377  -1.4376261  -0.6595934   0.31699112  0.6751458   0.99656415
  0.40565705 -1.0904227  -0.3513346  -0.66078615  1.1834346  -1.0899751
 -1.4925232  -0.30818892  1.4249563   0.06006899 -3.2386255   0.96192694
 -1.1045157   0.5540482  -1.5388466  -0.8721646   1.1221852   1.6488599
  0.44869688  1.1519432  -1.4588032  -0.04230021 -0.33113605  1.1316347
 -0.7425484  -0.11400439  0.37237874 -0.34573358  0.4140474  -0.04413145
  0.6157635  -1.0094129  -1.2208599  -0.7154122   0.9412035   0.9452426
 -0.0973389  -0.23566085  0.34300375 -0.95858365  0.8764276  -0.5669889
 -1.933235    0.22371146  1.6641699   1.3258857 ]