【pytorch实战1】用skip-gram训练词向量

首先将文本中的词汇用One-Hot Encoding表示，根据需要设置Word Vector维度，输入层变量个数及one-hot vector的维度（代码中为30000），隐藏层单元个数即为Word Vector维度，输出与输入维度相同，也是One-Hot Encoding。

Skip-gram原理如图

我们最终想要的是词库中单词的词向量表示，所以使用一层神经网络来实现Skip-gram算法，最后得到Word Embedding矩阵即可。

代码中使用负采样来进行模型训练，关于负采样，可以看看知乎上的这篇文章
word2vec中的负例采样为什么可以得到和softmax一样的效果

代码是跟着B站视频学习时敲的，用pytorch实现了训练，关于词向量算法原理的数学推导和assignment在学习CS224N的时候进行了实现，这里就直接用pytorch封装好的了。

USE_CUDA = torch.cuda.is_available()

# 为了保证实验结果可以复现，我们经常会把各种random seed固定在某一个值
random.seed(53113)
np.random.seed(53113)
torch.manual_seed(53113)
if USE_CUDA:
    torch.cuda.manual_seed(53113)

# 设定一些超参数
K = 100  # number of negative samples 负样本随机采样数量
C = 3  # nearby words threshold 指定周围三个单词进行预测
NUM_EPOCHS = 2  # The number of epochs of training 迭代轮数
MAX_VOCAB_SIZE = 30000  # the vocabulary size 词汇表多大
BATCH_SIZE = 128  # the batch size 每轮迭代1个batch的数量
LEARNING_RATE = 0.2  # the initial learning rate #学习率
EMBEDDING_SIZE = 100  # 词向量维度

LOG_FILE = "word-embedding.log"


# 文本转化成单词
def word_tokenize(text):
    return text.split()


with open("text8.train.txt", "r") as fin:  # 读入文件
    text = fin.read()

text = [w for w in word_tokenize(text.lower())]  # 分词
vocab = dict(Counter(text).most_common(MAX_VOCAB_SIZE - 1))  # 单词和对应出现次数的dict
vocab[""] = len(text) - np.sum(list(vocab.values()))  # 文本中其他单词都是不常用单词
idx_to_word = [word for word in vocab.keys()]  # 单词list
word_to_idx = {word: i for i, word in enumerate(idx_to_word)}  # {word：idx}
word_counts = np.array([count for count in vocab.values()], dtype=np.float32)  # 取出单词出现的次数，用于后面频率计算
word_freqs = word_counts / np.sum(word_counts)  # 得到每个单词的频率
word_freqs = word_freqs ** (3. / 4.)  # 论文里**3/4效果最好
word_freqs = word_freqs / np.sum(word_freqs)  # 重新计算所有word频率，用于负采样
VOCAB_SIZE = len(idx_to_word)


class WordEmbeddingDataSet(tud.Dataset):
    def __init__(self, text, word_to_idx, idx_to_word, word_freqs, word_counts):
        ''' text: a list of words, all text from the training dataset
                   word_to_idx: the dictionary from word to idx
                   idx_to_word: idx to word mapping
                   word_freq: the frequency of each word
                   word_counts: the word counts
               '''
        super(WordEmbeddingDataSet, self).__init__()
        
        '''
            text_encode是将整个文本（本例中也就是text8.txt）所有单词进行数字化编码，即把每个单词都用word_to_idx中该单词对应的idx来代替这个单词
            这样一段文本就变成了一串数字，比如i am a people就变成四个数字（i am a people这四个单词在字典中对应的idx，比如3 5 28 5333这样子）
        '''
        self.text_encode = [word_to_idx.get(t, VOCAB_SIZE - 1) for t in text]  # 把一段测试文本转化成了一个list，list里存放每个单词所对应的idx
        # 字典 get() 函数返回指定键的值（第一个参数），如果值不在字典中返回默认值（第二个参数）。
        # 取出text里每个单词word_to_idx字典里对应的索引,不在字典里返回""的索引
        # ""的索引=29999，get括号里第二个参数应该写word_to_idx[""]，不应该写VOCAB_SIZE-1，虽然数值一样。

        # print("text_encode")
        # print(self.text_encode)
        #
        # print("text_encode_len")
        # print(len(self.text_encode))
        #
        # print("word2idx")
        # print(word_to_idx)

        self.text_encode = torch.LongTensor(self.text_encode)
        # 变成tensor类型，这里变成longtensor，也可以torch.LongTensor(self.text_encode)

        self.word_to_idx = word_to_idx
        self.idx_to_word = idx_to_word
        self.word_freqs = word_freqs
        self.word_counts = word_counts

    def __len__(self):  # 数据集有多少个item
        # 魔法函数__len__
        ''' 返回整个数据集（所有单词）的长度
        '''
        return len(self.text_encode)

    # 这里的idx是text_encode里面对应的位置，关于text_encode的含义见init方法
    def __getitem__(self, idx):
        # 魔法函数__getitem__，这个函数跟普通函数不一样
        ''' 这个function返回以下数据用于训练
            - 中心词
            - 这个单词附近的(positive)单词
            - 随机采样的K个单词作为negative sample
        '''

        center_word = idx_to_word[idx]
        print("center word")
        print(center_word)
        # 中心词索引
        # 这里__getitem__函数是个迭代器，idx代表了所有的单词索引。

        pos_indices = list(range(idx - C, idx)) + list(range(idx + 1, idx + C + 1))
        print(pos_indices)

        # 周围词索引的索引，比如idx=0时。pos_indices = [-3, -2, -1, 1, 2, 3]
        pos_indices = [i % len(self.text_encode) for i in pos_indices]
        # range(idx+1, idx+C+1)超出词汇总数时，需要特别处理，取余数

        pos_words = self.text_encode[pos_indices]
        # 前后K个词的索引
        '''
           pos_indices:是相对中心词的位置的前后K个位置，比如中心词在句子中的位置是3，则前三和后三个分别是0，1，2，4，5，6
           pos_words:利用pos_indices和text_encode，找到前K和后K个词在单词表中对应的idx，这里相当于正例采样
        '''

        print("pos_words.shape:")
        print(pos_words.shape)
        print(len(pos_words))
        print(pos_words.size())

        neg_words = torch.multinomial(self.word_freqs, K * pos_words.shape[0], True)
        # 负例采样单词索引，torch.multinomial作用是对self.word_freqs做K * pos_words.shape[0]次取值，输出的是self.word_freqs对应的下标。
        # 取样方式采用有放回的采样，并且self.word_freqs数值越大，取样概率越大。
        # 每个正确的单词采样K个，pos_words.shape[0]是正确单词数量
        # print(neg_words)

        return center_word, pos_words, neg_words


dataset = WordEmbeddingDataSet(text, word_to_idx, idx_to_word, word_freqs, word_counts)
list(dataset)  # 可以尝试打印下center_word, pos_words, neg_words看看

dataloader = tud.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)


class EmbeddingModel(nn.Module):
    def __init__(self, vocab_size, embed_size):
        super(EmbeddingModel, self).__init__()
        self.vocab_size = vocab_size  # 30000
        self.embed_size = embed_size  # 100

        initrange = 0.5 / self.embed_size

        # out_embed和in_embed是两个Embedding类的对象
        self.out_embed = nn.Embedding(self.vocab_size, self.embed_size, sparse=False)  # sparse稀疏
        # 模型输出nn.Embeding(30000*100)
        self.out_embed.weight.data.uniform_(-initrange, initrange)  # 权重初始化
        # uniform(x, y) 方法将随机生成下一个实数，它在 [x, y] 范围内。

        self.in_embed = nn.Embedding(self.vocab_size, self.embed_size, sparse=False)  # sparse稀疏
        # 模型输入nn.Embeding(30000*100)
        self.in_embed.weight.data.uniform_(-initrange, initrange)  # 权重初始化

    def forward(self, input_labels, pos_labels, neg_labels):
        '''
        :param input_labels: 中心词 [batch_size]
        :param pos_labels: 中心词周围 context window 出现过的单词 [batch_size * (window_size * 2)]
        :param neg_labels: 中心词周围没有出现过的单词，从 negative sampling 得到 [batch_size, (window_size * 2 * K)]
        :return: loss [batch_size]
        '''

        print("input_labels:")
        print(input_labels)
        batch_size = input_labels.size(0)  # 中心词数目

        #用in_embed和input_labels做embedding操作
        input_embedding = self.in_embed(input_labels)  # 128(batch_size) * 100(embed_size)
        # 这里估计进行了运算：（128,30000）*（30000,100）= 128(B) * 100 (embed_size)

        pos_embedding = self.out_embed(pos_labels)  # B * (2*C) * embed_size
        # 同上，增加了维度(2*C)，表示一个batch有B组周围词单词，一组周围词有(2*C)个单词，每个单词有embed_size个维度。

        neg_embedding = self.out_embed(neg_labels)  # B * (2*C * K) * embed_size
        # 同上，增加了维度(2*C*K)

        # torch.bmm()为batch间的矩阵相乘（b,n.m)*(b,m,p)=(b,n,p)
        log_pos = torch.bmm(pos_embedding, input_embedding.unsqueeze(2)).squeeze()  # B * (2*C)
        # 这句对应算法里中心词与前后K个词进行向量内积运算

        log_neg = torch.bmm(neg_embedding, -input_embedding.unsqueeze(2)).squeeze()  # B * (2*C*K)
        # 这句对应算法里中心词与负采样进行向量内积运算再取相反数
        # unsqueeze(2)指定位置升维，.squeeze()压缩维度。

        # 论文里的损失计算公式
        log_pos = F.logsigmoid(log_pos).sum(1)
        log_neg = F.logsigmoid(log_neg).sum(1)  # 对第二维求和了，所以log_pos和log_neg都是[batch_size]
        loss = log_pos + log_neg

        return -loss

    def input_embeddings(self):  # 取出self.in_embed数据参数
        return self.in_embed.weight.data.cpu().numpy()

model = EmbeddingModel(VOCAB_SIZE, EMBEDDING_SIZE)
#得到model，有参数，有loss，可以优化了

if USE_CUDA:
    model = model.cuda()

optimizer = torch.optim.SGD(model.parameters(), lr = LEARNING_RATE)

for e in range(NUM_EPOCHS):
    for i, (input_labels, pos_labels, neg_labels) in enumerate(dataloader):

        input_labels = input_labels.long()
        pos_labels = pos_labels.long()
        neg_labels = neg_labels.long()

        if USE_CUDA:
            input_labels = input_labels.cuda()
            pos_labels = pos_labels.cuda()
            neg_labels = neg_labels.cuda()

            optimizer.zero_grad()

            loss = model(input_labels, pos_labels, neg_labels).mean()

            loss.backward()
            optimizer.step()