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)
syn0
にnumpy.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