機械学習・自然言語処理の勉強メモ

学んだことのメモやまとめ

Pytorch:Embeddingに学習済みの重みを使う

やりたいこと



事前にWord2Vecなどで学習した分散表現をネットワークの重みとして用いる。

結論としては、Embedding層の重みにテンソル型に変換した分散表現行列をセットするだけで良かった。

備忘録としてコードを残しておく。
(ここでは、Googleの学習済みの分散表現ベクトルを用いた。)
これのハンドリングが良く分からったので、いったんgensimで読み込んだ。

import torch
import torch.nn as nn
import numpy as np
import gensim

model_dir = './GoogleNews-vectors-negative300.bin'
model = gensim.models.KeyedVectors.load_word2vec_format(model_dir, binary=True)

syn0numpy.ndarray型の重みが格納されている。

word_vectors = model.wv
weights = word_vectors.syn0

type(weights)  # numpy.ndarray
weights.shape  # (3000000, 300)

続いて以下のようにしてEmbedding層を定義する。

vocab_size = weights.shape[0]
embedding_dim = weights.shape[1]
embed = nn.Embedding(vocab_size, embedding_dim)

最後に以下で学習済みの重みをセットする。

embed.weight = nn.Parameter(torch.from_numpy(weights))

尚、重みを固定したい場合、以下のようにしておく。

embed.weight.requires_grad = False

1点注意であるが、分散表現行列の各行は単語のベクトルであり、「行番号が単語のインデックス」と対応している。この暗黙的なマッピングが無いと、無関係な単語に対する重みをセットしていることになり、学習済みの重みをセットする意味がなくなる。

今回の内容はここのディスカッションの内容を整理したものである。
discuss.pytorch.org