word2vecのCBOWモデルをPytorchで実装してみた

説明

『ゼロから作るDeep-Learning2』という本を読み進めながら勉強をしている過程でword2vecのCBOWというモデルに関する解説・実装がありました。
仕組みを深く理解するために自分で実装したいと思ったのですが本に書いてあるnumpyのみを使った実装を丸写しするだけでは面白くないと思い、並列して勉強中だったPytorchを使ってモデルを作ることにしました。
一部のクラスや関数を『ゼロから作るDeep-Learning2』内から引用して使用しています。
また詳しい仕組みに関しても同書に非常にわかりやすく説明してあるのでぜひ参考にしてください。

実行環境

• Intel(R) Core(TM) i7-8700 6コア12スレッド
• メモリ16GB
• Windows10 64bit
• Docker Desktop for Windows
• 使用DockerImage pytorch/pytorch : https://hub.docker.com/r/pytorch/pytorch
• Python 3.7.13
• JupyterLab

CBOWとは？

CBOWモデルとは、Word2vecと呼ばれる単語の分散表現を生成するために使用される種類のモデル群のうちの一つです。
このモデルの目標はそれぞれの単語に対しその意味を表現する固定長のベクトルを生成することにあります。
今回は1万種類の単語を含む100万語の文章データから、1万個の単語それぞれの意味に対応した単語ベクトルを得ることを目標にしています。
具体的な仕組みはコードを見ていただくか、『ゼロから作るDeep-Learning2』をお読みください
また以下の記事も非常によく説明しており参考になります
https://qiita.com/g-k/items/69afa87c73654af49d36

コード

モジュール

modules

import sys
import os
sys.path.append('..')
import pickle
import urllib.request #urllib3 Version: 1.26.8

import numpy as np #Version: 1.21.5
import torch #Version: 1.12.0
import torch.nn as nn
import torch.utils.data as data
import torch.optim as optim
import torch.autograd as autograd
from torchvision import datasets, transforms #Version: 0.13.0

from tqdm import tqdm
#今回の自分の環境などjupyter環境の場合以下のコード
#from tqdm.notebook import tqdm #Version: 4.63.0

データの用意

以下の前処理コードはサンプルコードからの引用です。データをダウンロードしコーパス及び数値と単語を結びつける辞書を作成します。

detaset

# [出典]:サンプルコード dataset/ptb.py

url_base = 'https://raw.githubusercontent.com/tomsercu/lstm/master/data/'
key_file = {
    'train':'ptb.train.txt',
    'test':'ptb.test.txt',
    'valid':'ptb.valid.txt'
}
save_file = {
    'train':'ptb.train.npy',
    'test':'ptb.test.npy',
    'valid':'ptb.valid.npy'
}
vocab_file = 'ptb.vocab.pkl'

#dataset_dir = os.path.dirname(os.path.abspath(__file__))
dataset_dir = "/"

def _download(file_name):
    file_path = dataset_dir + '/' + file_name
    if os.path.exists(file_path):
        return

    print('Downloading ' + file_name + ' ... ')

    try:
        urllib.request.urlretrieve(url_base + file_name, file_path)
    except urllib.error.URLError:
        import ssl
        ssl._create_default_https_context = ssl._create_unverified_context
        urllib.request.urlretrieve(url_base + file_name, file_path)

    print('Done')


def load_vocab():
    vocab_path = dataset_dir + '/' + vocab_file

    if os.path.exists(vocab_path):
        with open(vocab_path, 'rb') as f:
            word_to_id, id_to_word = pickle.load(f)
        return word_to_id, id_to_word

    word_to_id = {}
    id_to_word = {}
    data_type = 'train'
    file_name = key_file[data_type]
    file_path = dataset_dir + '/' + file_name

    _download(file_name)

    words = open(file_path).read().replace('\n', '<eos>').strip().split()

    for i, word in enumerate(words):
        if word not in word_to_id:
            tmp_id = len(word_to_id)
            word_to_id[word] = tmp_id
            id_to_word[tmp_id] = word

    with open(vocab_path, 'wb') as f:
        pickle.dump((word_to_id, id_to_word), f)

    return word_to_id, id_to_word


def load_data(data_type='train'):
    '''
        :param data_type: データの種類：'train' or 'test' or 'valid (val)'
        :return:
    '''
    if data_type == 'val': data_type = 'valid'
    save_path = dataset_dir + '/' + save_file[data_type]

    word_to_id, id_to_word = load_vocab()

    if os.path.exists(save_path):
        corpus = np.load(save_path)
        return corpus, word_to_id, id_to_word

    file_name = key_file[data_type]
    file_path = dataset_dir + '/' + file_name
    _download(file_name)

    words = open(file_path).read().replace('\n', '<eos>').strip().split()
    corpus = np.array([word_to_id[w] for w in words])

    np.save(save_path, corpus)
    return corpus, word_to_id, id_to_word

create_contexts_target関数です。コーパスから学習の入力として使うコンテキスト（ターゲットの単語の左右の単語）とターゲットの単語の組を作成します。

create_contexts_target

# [出典]:サンプルコード common/util.py
def create_contexts_target(corpus, window_size=1):
    '''コンテキストとターゲットの作成

    :param corpus: コーパス（単語IDのリスト）
    :param window_size: ウィンドウサイズ（ウィンドウサイズが1のときは、単語の左右1単語がコンテキスト）
    :return:
    '''
    target = corpus[window_size:-window_size]
    contexts = []

    for idx in range(window_size, len(corpus)-window_size):
        cs = []
        for t in range(-window_size, window_size + 1):
            if t == 0:
                continue
            cs.append(corpus[idx + t])
        contexts.append(cs)

    return np.array(contexts), np.array(target)

class CBOWDataset(data.Dataset):
    def __init__(self, window_size, corpus):
        self.corpus = corpus
        self.contexts, self.targets = create_contexts_target(self.corpus, window_size=window_size)

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

    def __getitem__(self, index):
        context =  self.contexts[index]
        target = self.targets[index]
        return torch.tensor(context), torch.tensor(target), self.corpus

以下は学習を効率化するために「Negative Sampling」を行うクラスです。

UnigramSampler

#[出典]:サンプルコード ch04/negative_sampling_layer.py
import collections
GPU=True
class UnigramSampler:
    def __init__(self, corpus, power, sample_size):
        self.sample_size = sample_size
        self.vocab_size = None
        self.word_p = None

        counts = collections.Counter()
        for word_id in corpus:
            counts[word_id] += 1

        vocab_size = len(counts)
        self.vocab_size = vocab_size

        self.word_p = np.zeros(vocab_size)
        for i in range(vocab_size):
            self.word_p[i] = counts[i]

        self.word_p = np.power(self.word_p, power)
        self.word_p /= np.sum(self.word_p)

    def get_negative_sample(self, target):
        batch_size = target.shape[0]

        if not GPU:
            negative_sample = np.zeros((batch_size, self.sample_size), dtype=np.int32)

            for i in range(batch_size):
                p = self.word_p.copy()
                target_idx = target[i]
                p[target_idx] = 0
                p /= p.sum()
                negative_sample[i, :] = np.random.choice(self.vocab_size, size=self.sample_size, replace=False, p=p)
        else:
            # GPU(cupy）で計算するときは、速度を優先
            # 負例にターゲットが含まれるケースがある
            negative_sample = np.random.choice(self.vocab_size, size=(batch_size, self.sample_size),
                                               replace=True, p=self.word_p)

        return negative_sample

Dataset

ここまでが原書から引用した前処理コードとなります。ここからpytorchによる学習を行うためにDatasetを作成していきます。

CBOWDataset

class CBOWDataset(data.Dataset):
    """Pytorch用にDatasetを作成するクラス
    
    Parameters:
    ----------
    window_size : int
        ターゲットの単語の左右に何個のの単語を使用するか
    corpus : list
        コーパス

    Returns:
    ----------
    テンソルのコンテキスト、ターゲット、入力そのままのcorpus
    torch.tensor(context), torch.tensor(target), self.corpus
    """
    def __init__(self, window_size, corpus):
        self.corpus = corpus
        self.contexts, self.targets = create_contexts_target(self.corpus, window_size=window_size)

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

    def __getitem__(self, index):
        context =  self.contexts[index]
        target = self.targets[index]
        return torch.tensor(context), torch.tensor(target), self.corpus

以下でdatasetを作成します

dataset

#コーパスを用意します。validは今回検証を行わないので必要ありません。
corpus_train, word_to_id, id_to_word = load_data('train')
#corpus_valid, word_to_id, id_to_word = load_data('valid')
window_size = 5
dataset_train = CBOWDataset(window_size, corpus_train)

DataLoader

次にミニパッチごとにデータを取得するDataLoaderを作成します。
ミニパッチを作るだけでなく学習の効率化のためにNegative Samplingを行う必要があり、DataLoaderに入力する関数od_collate_fn内で行っています。入力されたDatasetからミニパッチの数だけデータを受け取り、
(コンテキスト、ターゲット、コーパス)のデータから、ミニパッチの分だけ
(コンテキスト, x(ターゲットとNegative Samplingの結果10個), y(y[0]=1でそれ以外10個がゼロ。xに対応する正解データ))
というデータを作成します

dataloader

def od_collate_fn(batch):
    contexts = []
    targets = []

    sample_size = 10
    power = 0.75
    corpus_use = batch[0][2]
    sampler = UnigramSampler(corpus_use, power, sample_size)

    for sample in batch:
        contexts.append(sample[0])
        targets.append(sample[1])
    negative_sample = torch.tensor(sampler.get_negative_sample(np.array(targets)))
    contexts = torch.stack(contexts, dim=0)
    targets = torch.tensor(targets)

    
    x = torch.cat([targets.view(-1, 1),negative_sample],dim=1)
    y = []
    for dim0 in range(x.size()[0]):
        y.append(torch.where(x[dim0] == x[dim0][0], torch.ones_like(x[dim0]), torch.zeros_like(x[dim0])).view(1, -1))
    y = torch.cat(y,dim=0)
    #y = y/(y.sum(dim=1).view(-1,1))
    return contexts, x, y

batch_size = 100

dataloader_train = torch.utils.data.DataLoader(
    dataset_train, #CBOWDatasetで作成
    batch_size=batch_size,
    shuffle=True,
    collate_fn = od_collate_fn
)

モデルの作成

肝心のCBOWモデルを作成します
これは重みのパラメーター行列から入力のindexに対応する行を抜き出してくるEmbeddingレイヤーです。
モデルを作るうえでの基盤になります
Pytorchはautogradで自動的にbackwardが計算されるためモデルのコードが非常にシンプルに書けるのが魅力です。

Embedding

class Embedding(nn.Module):  # nn.Moduleを継承する
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.W_in = nn.Parameter(torch.tensor(rng.uniform(
                        low=-np.sqrt(6/in_dim),
                        high=np.sqrt(6/in_dim),
                        size=(in_dim, out_dim)
                    ).astype('float32')))
    def forward(self, idx):
        return self.W_in[idx]

Embeddingレイヤーをもとに第1層のレイヤであるEmbeddingInレイヤーを作ります
入力のindexに対応する単語ベクトルすべてを加算し個数で割ることで縦に平均を取ります。

EmbeddingIn

class EmbeddingIn(nn.Module):  # nn.Moduleを継承する
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.embed = Embedding(in_dim, out_dim)

    def forward(self, idx):#idxは複数。平均をとる
        return self.embed.forward(idx).sum(idx.dim()-1)/idx.shape[idx.dim()-1]

2層目のレイヤーであるEmbeddingDotレイヤーを作ります。
入力された行列に対し2層目のパラメーター行列から単語ベクトルを抜き出してきて掛け合わせる操作なのですが、forwardの実装が少しトリッキーになっています。

EmbeddingDot

class EmbeddingDot(nn.Module):  # nn.Moduleを継承する
    def __init__(self, in_dim, out_dim):  # __init__をoverride
        super().__init__()
        self.embed = Embedding(in_dim, out_dim)

    def forward(self, h, idx):  # forwardをoverride
        return (self.embed.forward(idx).permute(1, 0, 2)*h).sum(2).permute(1, 0)

以上のレイヤーをすべて組み合わせてCBOWモデルを作ります
最終的な出力が二値分類問題に合うようにSigmoid関数を適用しています

CBOW

class CBOW(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.Embedding1 = EmbeddingIn(in_dim, out_dim)
        self.Embedding2 = EmbeddingDot(in_dim, out_dim)

    def forward(self, contexts, x):
        h = self.Embedding1(contexts)
        y = self.Embedding2(h, x)
        y = torch.sigmoid(y)
        return y

学習

学習の準備です。
ハイパーパラメーター、モデル、最適化手法を設定しています
最適化手法はAdamを使用しています。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

in_dim = len(word_to_id)
hid_dim = 100
n_epochs = 10
lr = 0.001


cbow = CBOW(in_dim, hid_dim).to(device)

optimizer = optim.Adam(cbow.parameters(), lr=lr)

実際に学習を行います。環境によりますが12時間ほど終了にかかります
1エポックごとにパラメーターのバックアップを取っています。

training

for epoch in tqdm(range(n_epochs)):
    losses_train = []
    losses_valid = []
    train_num = 0
    train_true_num = 0
    valid_num = 0
    valid_true_num = 0

    cbow.train()  # 訓練時には勾配を計算するtrainモードにする
    for i, (x1,x2, t) in enumerate(tqdm(dataloader_train, leave=False)):
        # 勾配の初期化
        cbow.zero_grad()

        # テンソルをGPUに移動
        x1 = x1.to(device)
        x2 = x2.to(device)
        t = t.to(device)
        
        # 順伝播
        y = cbow.forward(x1, x2)

        # 誤差の計算(二値分類用クロスエントロピー誤差関数)
        loss = (-(t*torch.log(y))-(1-t)*(torch.log(1-y))).sum(axis=1).mean()

        # 誤差の逆伝播
        optimizer.zero_grad()
        loss.backward()

        # パラメータの更新
        optimizer.step()

        pred = y.argmax(1)
        t_ = t.argmax(1)
        losses_train.append(loss.tolist())
        acc = torch.where(t_.to("cpu") - pred.to("cpu") == 0, torch.ones_like(t_.to("cpu")), torch.zeros_like(t_.to("cpu")))
        train_num += acc.size()[0]
        train_true_num += acc.sum().item()
    torch.save(cbow.state_dict(), 'model{}.pth'.format(epoch))#パラメーターのバックアップ
    print('EPOCH: {}, Train [Loss: {:.3f}, Accuracy: {:.3f}]'.format(
        epoch,
        np.mean(losses_train),
        train_true_num/train_num,
    ))
    

結果の確認

結果の確認にはまた『ゼロから作るDeep-Learning2』内で使われていたコードを使おうと思います

similarity

#[出典]: サンプルコード common/util.py

def cos_similarity(x, y, eps=1e-8):
    '''コサイン類似度の算出

    :param x: ベクトル
    :param y: ベクトル
    :param eps: ”0割り”防止のための微小値
    :return:
    '''
    nx = x / (np.sqrt(np.sum(x ** 2)) + eps)
    ny = y / (np.sqrt(np.sum(y ** 2)) + eps)
    return np.dot(nx, ny)
def most_similar(query, word_to_id, id_to_word, word_matrix, top=5):
    '''類似単語の検索

    :param query: クエリ（テキスト）
    :param word_to_id: 単語から単語IDへのディクショナリ
    :param id_to_word: 単語IDから単語へのディクショナリ
    :param word_matrix: 単語ベクトルをまとめた行列。各行に対応する単語のベクトルが格納されていることを想定する
    :param top: 上位何位まで表示するか
    '''
    if query not in word_to_id:
        print('%s is not found' % query)
        return

    print('\n[query] ' + query)
    query_id = word_to_id[query]
    query_vec = word_matrix[query_id]

    vocab_size = len(id_to_word)

    similarity = np.zeros(vocab_size)
    for i in range(vocab_size):
        similarity[i] = cos_similarity(word_matrix[i], query_vec)

    count = 0
    for i in (-1 * similarity).argsort():
        if id_to_word[i] == query:
            continue
        print(' %s: %s' % (id_to_word[i], similarity[i]))

        count += 1
        if count >= top:
            return

pytorchのモデルからパラメータを引き出しnumpy行列としたうえでpickleファイルとして保存しておきます

pickle

word_vecs = cbow.state_dict()["Embedding1.embed.W_in"].to('cpu').detach().numpy().copy()
params = {} 
params['word_vecs'] = word_vecs.astype(np.float16) 
params['word_to_id'] = word_to_id 
params['id_to_word'] = id_to_word
pkl_file = 'cbow_params.pkl' 
with open(pkl_file, 'wb') as f: 
    pickle.dump(params, f, -1)

保存された行列のpickleファイルを読み込んで、試しにいくつかの単語に対して近いベクトルを持つ単語を5つ表示してみます。

test

import sys 
import pickle 
pkl_file = 'cbow_params.pkl' 
with open(pkl_file, 'rb') as f: 
    params = pickle.load(f) 
    word_vecs = params['word_vecs'] 
    word_to_id = params['word_to_id'] 
    id_to_word = params['id_to_word'] 
querys = ['you', 'year', 'car', 'toyota', 'high'] 
for query in querys: 
    most_similar(query, word_to_id, id_to_word, word_vecs, top=5)

学習に成功していることが分かりますね。似たような使われ方をする単語が同じようなベクトルになるようなモデルなので、highの最上位に対義語であるlowが入ってしまっています。このような現象を解決するためにはより発展的なモデルを使う必要があります。

[query] you
 we: 0.75634765625
 i: 0.712890625
 they: 0.63720703125
 your: 0.62646484375
 anybody: 0.5888671875

[query] year
 month: 0.8447265625
 week: 0.759765625
 summer: 0.75048828125
 spring: 0.7080078125
 decade: 0.6943359375

[query] car
 auto: 0.626953125
 luxury: 0.62158203125
 window: 0.580078125
 cars: 0.572265625
 scorpio: 0.55322265625

[query] toyota
 seita: 0.671875
 nissan: 0.6416015625
 minicomputers: 0.61865234375
 chevrolet: 0.6162109375
 ford: 0.61181640625

[query] high
 low: 0.76318359375
 steady: 0.485107421875
 peak: 0.436767578125
 taxable: 0.429931640625
 refined: 0.4267578125