今日も窓辺でプログラム

外資系企業勤めのエンジニアが勉強した内容をまとめておくブログ

TensorFlowでword2vecを使って単語ベクトルを学習する

今回やること

前回の記事で、RNNを使って文字レベルの言語モデルを実装しました。
シンプルなRNNで文字レベルの言語モデルをTensorFlowで実装してみる - 今日も窓辺でプログラム

英語を対象にしたので入出力の次元は26文字+スペースの27次元で済んだのですが、単語レベルの言語モデルを実装しようとすると次元は非常大きなものになってしまいます。
それを解決するのが皆さんご存知word2vecで、TensorFlowではword2vecが簡単に利用できるような機能が提供されているようなので、今回はそれを試してみたいと思います。

実装には、以下のコードを大いに参考にしています。
tensorflow/word2vec_basic.py at master · tensorflow/tensorflow · GitHub

コードの場所

今回のコードも、GitHubに上げておきます。下記のcorpus.pyという1つのファイルに収まってます。
github.com

word2vecの基礎知識

まずはword2vecに概要を把握することから始めないといけません。絵や数式を書いたりするのは得意ではないので…丁寧に説明されている記事をいくつか紹介しておきます。

絵で理解するWord2vecの仕組み - Qiita
絵を多用し、数式を使わずに、word2vecの仕組みが丁寧に説明されています。

(原文英語) Vector Representations of Words  |  TensorFlow
(日本語訳) TensorFlowチュートリアル - 単語のベクトル表現(翻訳) - Qiita
TensorFlow公式サイトの解説です。数式も出てきますが、図も多用して説明されています。

コーパスをone-hot表現に変換

上述の解説記事を読んでいただくと分かるかと思いますが、まずはコーパスをone-hot表現に変換する必要があります。
one-hot表現は、文字レベルの言語モデルでも使いましたが、1つの要素だけが1、他が0のベクトルでした。
単語"apple"の単語IDが0だった場合は、[1, 0, 0, ..., 0]というようなベクトルになります。

今回実装したコードを、最初のパラメータの定義とともにコードを載せておきます。

class Corpus:
    def __init__(self):
        # パラメータ
        self.embedding_size = 100
        self.batch_size = 8
        self.num_skips = 2
        self.skip_window = 1
        self.num_epochs = 1000
        self.learning_rate = 0.1

        self.current_index = 0
        self.words = []

        self.dictionary = {} # コーパスの単語と単語ID
        self.final_embeddings = None # 最終的なベクトル表現

    def build_dataset(self):
        new_word_id = 0
        self.words = []
        self.dictionary = {}

        # コーパスとなるファイルたちを順に読み込む
        for filename in glob.glob("./corpus/*.txt"):
            with open(filename, "r", encoding="utf-8") as f:
                # 簡単な前処理をして、小文字表記の単語がスペースで区切られた状態にする。
                text = f.read()
                text = text.lower().replace("\n", " ")
                text = re.sub(r"[^a-z '\-]", "", text)
                text = re.sub(r"[ ]+", " ", text)

                for word in text.split():
                    # 新しい単語はdictionaryに登録
                    if word.startswith("-"): continue # 簡単な前処理。"-"から始まる単語は無視。
                    if word not in self.dictionary:
                        self.dictionary[word] = new_word_id
                        new_word_id += 1
                    self.words.append(self.dictionary[word])

        # 本当は、出現頻度の低い単語を「未知語」としてひとまとめにしたほうがよいが、
        # 今回はその処理はしないことにします。。
        self.vocabulary_size = new_word_id
        print("# of distinct words:", new_word_id)
        print("# of total words:", len(self.words))

Skip-gram学習用のバッチを作成

今回はSkip-gramを使って単語ベクトルを学習させていきます。Skip-gram自体の詳細については、他サイトの記事に解説をお任せします。

日本語だとこの記事なんかどうでしょう。
Word2Vec のニューラルネットワーク学習過程を理解する · けんごのお屋敷

先述のTensorFlow公式の記事でもSkip-gramは解説されています。
(原文英語) Vector Representations of Words  |  TensorFlow
(日本語訳) TensorFlowチュートリアル - 単語のベクトル表現(翻訳) - Qiita

では実際に、注目している単語とその周辺の単語のペアを教師データとして出力する部分を書いていきます。
学習を回す部分で使いやすいように、Pythonのジェネレータを使って、パラメータ self.batch_size で決められたサイズのバッチを生成するようなコードにしたいと思います。

この部分の実装についても、TensorFlow公式のサンプルを参考にしています。(というか、ほぼそのままです)

    # skip-gramのバッチを作成
    def generate_batch(self):
        assert self.batch_size % self.num_skips == 0
        assert self.num_skips <= 2 * self.skip_window

        self.current_index = 0
        batch = np.ndarray(shape=(self.batch_size), dtype=np.int32) # 注目してる単語
        labels = np.ndarray(shape=(self.batch_size, 1), dtype=np.int32) # その周辺の単語

        # 次の処理範囲分のテキストがなかったらイテレーション終了
        span = 2 * self.skip_window + 1
        if self.current_index + span >= len(self.words):
            raise StopIteration

        # 今処理している範囲をbufferとして保持する
        buffer = collections.deque(maxlen=span)
        for _ in range(span):
            buffer.append(self.words[self.current_index])
            self.current_index += 1

        # バッチサイズごとにyeildで結果を返すためのループ
        for _ in range(len(self.words) // self.batch_size):
            # 注目している単語をずらすためのループ
            for i in range(self.batch_size // self.num_skips):
                target = self.skip_window
                targets_to_avoid = [self.skip_window]
                # 注目している単語の周辺の単語用のループ
                for j in range(self.num_skips):
                    while target in targets_to_avoid:
                        target = random.randint(0, span - 1)
                    targets_to_avoid.append(target)
                    batch[i * self.num_skips + j] = buffer[self.skip_window]
                    labels[i * self.num_skips + j, 0] = buffer[target]

                # 今注目している単語は処理し終えたので、処理範囲をずらす
                buffer.append(self.words[self.current_index])
                self.current_index += 1
                if self.current_index >= len(self.words):
                    raise StopIteration
            yield batch, labels
        raise StopIteration

学習部分の実装

学習部分はも先ほどのサンプルコードを参考にしています。
またコードを全部紹介しますが、先ほど紹介した記事などに載っているグラフを構築して、先ほど用意したバッチデータを使って学習を進めていくのみです。

今回は学習したネットワークのパラメータを保存するよりも、単語IDと学習済みの単語ベクトルを保存しておいたほうが後の役に立つので、関数の最後で単語IDと各単語のベクトル表現をそれぞれファイルに保存しています。

    def train(self):
        # 単語ベクトルの変数を用意
        embeddings = tf.Variable(
            tf.random_uniform([self.vocabulary_size, self.embedding_size], -1.0, 1.0))

        # NCE用の変数
        nce_weights = tf.Variable(
            tf.truncated_normal([self.vocabulary_size, self.embedding_size],
                                stddev=1.0 / math.sqrt(self.embedding_size)))
        nce_biases = tf.Variable(tf.zeros([self.vocabulary_size]))

        # 教師データ
        train_inputs = tf.placeholder(tf.int32, shape=[self.batch_size])
        train_labels = tf.placeholder(tf.int32, shape=[self.batch_size, 1])

        # 損失関数
        embed = tf.nn.embedding_lookup(embeddings, train_inputs)
        loss = tf.reduce_mean(
            tf.nn.nce_loss(nce_weights, nce_biases, embed, train_labels, self.batch_size // 2, self.vocabulary_size)
        )

        # 最適化
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=self.learning_rate).minimize(loss)

        norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keep_dims=True))
        normalized_embeddings = embeddings / norm

        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())

            # 決められた回数エポックを回す
            for epoch in range(self.num_epochs):
                epoch_loss = 0
                # generate_batch()で得られたバッチに対して学習を進める
                for batch_x, batch_y in self.generate_batch():
                    _, loss_value = sess.run([optimizer, loss], feed_dict={train_inputs: batch_x, train_labels: batch_y})
                    epoch_loss += loss_value

                print("Epoch", epoch, "completed out of", self.num_epochs, "-- loss:", epoch_loss)

            # 一応モデルを保存
            saver = tf.train.Saver()
            saver.save(sess, "./corpus/model/blog.ckpt")

            # 学習済みの単語ベクトル
            self.final_embeddings = normalized_embeddings.eval() # <class 'numpy.ndarray'>


        # 単語IDと学習済みの単語ベクトルを保存
        with open("./corpus/model/blog.dic", "wb") as f:
            pickle.dump(self.dictionary, f)
        print("Dictionary was saved to", "./corpus/model/blog.dic")
        np.save("./corpus/model/blog.npy", self.final_embeddings)
        print("Embeddings were saved to", "./corpus/model/blog.npy/")

小さなコーパスで学習させてみる

使用するコーパス

文字レベルの言語モデルを作った時と同様、American National CorpusMASC (500K) – data onlyというコーパスの中の書き言葉のblogカテゴリの記事だけを使います。
単語数が27391、単語の種類は5423となっています。

コードが動いているか確認ということで、このコーパスに対して100次元のベクトル表現を学習させてみたいと思います。Skip-gramのウインドウは前後1単語で、学習は1000エポック回すことにしました。

結果を可視化してみる

100次元のベクトルを2次元に落として可視化したものが下記になります。(一部の単語のみをプロットしています)
f:id:kanohk:20161224182506p:plain

コーパスが小さいので実用レベルの表現にはなっていないとは思いますが、いい感じに単語が散らばっていることが確認できます。

おわりに

今回は、単語のベクトル表現を学習するコードをとりあえず実装しました。
小さいコーパスに対してコードを走らせて、それっぽいものが学習できているらしい、ということは確認しました。

次回は、この学習したベクトル表現が意味のあるものなのかを確かめられたらいいなと思います。



次回記事

TensorBoardにEmbedding Visualizationという機能があるようなので、それを使って可視化をしてみました。
少ないコード変更で使えるので、ぜひチェックしてみてください。
www.madopro.net