pytorchを使ってword2vecを実装してみたので、実装しながら陥ったバグや大変だったとこをまとめます。

一覧

modelの保存方法
最大化を最小化を間違えた
negative sampling用のtable作成
初期値

modelの保存方法

そのままモデルを保存するとGPU情報が含まれてしまい、ロードするときに自動的にGPUに乗ってしまう問題。
https://qiita.com/jyori112/items/aad5703c1537c0139edb
↑でまとめました。

最大化と最小化を間違えた

word2vec (CBoW)の最適化関数は

\mathcal{O}=\log\sigma(w\cdot h)+\sum_{p=1}^P\log\sigma(-w_p\cdot h)

ですが、これは最大化。他のDeep Learning用のライブラリと同様にpytorchでも普通は最小化をするので、-1を掛けて

\mathcal{O}=-\log\sigma(w\cdot h)-\sum_{p=1}^P\log\sigma(-w_p\cdot h)

を最小化するようにプログラムを組む必要がある。

negative sampling用のtable作成

各ステップごとにnegative sampleを確率的に選ぶ必要がありますが、np.random.choiceなどを使うと非常に遅い。
そこで、事前に用意したtableからランダムにサンプルする。

実装するのが意外と大変で数時間溶ける。

いろいろ工夫して最終的にできたものがこちら。

build_neg_table

def build_neg_table(self, power=3/4, table_size=10000000):
    powered_sum = sum(count ** power for wid, count in self.id2count.items())
    table = np.zeros(shape=(table_size, ),dtype=np.int32)

    idx = 0
    accum = 0.0

    for wid, count in self.id2count.items():
        freq = (count**power) / powered_sum
        accum += freq * table_size
        end_idx = int(accum)
        table[idx: int(accum)] = wid
        idx = end_idx

    return table

初期値

word2vecの学習では初期値が結果に大きく影響を与える。
pytorchのデフォルトの初期化では値が大きすぎて学習できないので、自前で実装。

# Make embeddings
self.center_embed = nn.Embedding(n_vocab, dim, padding_idx=0)
self.context_embed = nn.Embedding(n_vocab, dim, padding_idx=0)

# Initialize Embeddings
self.center_embed.weight.data.uniform_(-0.5 / dim, 0.5 / dim)
self.context_embed.weight.data.zero_()

最後に

ニューラル系の実装のデバッグはどうやったらいいのか未だにわからない。

ソースコードは
https://github.com/jyori112/pytorch-word2vec
においてあります。