RNN/LSTMを使った言語モデルをTensorFlowで実装してみる

はじめに

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

今回は、前回のコードを少しだけいじって、単語レベルの言語モデルを実装します。また、RNNのセルも、単純なものからLSTMに切り替えてみます。

今回やること

今回やることは、前回の記事の知識を前提とします。前回の記事をまだ読んでいない方は、先にそちらに目を通していただけるとこの記事が読みやすいかもしれません。

作りたい言語モデル

単語の列 $w_0, w_1, \dots, w_{n-1}$ というコンテキストが与えられたときに、次に来る単語 $w_n$ の確率を出力するモデルを作成します。
つまり、次のような確率がモデルの出力です。
$P(w_n | w_0, w_1, \cdots, w_{n-1})$

ニューラルネットの構成

入力はn個の単語、出力は次に来る単語です。
入出力の単語は語彙数の次元を持つベクトルのone-hot表現で扱います。つまり、入力層と出力層は語彙数の分だけノードが並んだ形になります。

隠れ層は、LSTMのユニットの層を1層だけ用意します。詳しくは後述しますが、個数は今回は30個並べます。
言語モデルというと入力層と隠れ層の中にEmbedding層を導入しているものもありますが、今回はいきなりLSTMにつないでしまいましょう。

使用するコーパス

American National CorpusのMASC (500K) – data onlyの一部を使用します。（文字レベルの時と同じサイトです）
今回は、少しフォーマルな内容の"newspaper_newswire"というフォルダの中にあるファイル(161KB)を使用します。

実装

コード全体

まず、コード全体はいつものようにGitHubにあげておくので必要に応じて参照していただければと思います。
NNLM/wordmodel.py at ac6f35b8dd18895f934bab4216ea71b596df1ab9 · kanoh-k/NNLM · GitHub

コーパスから辞書を作る

まずはコーパス全体を読み込み、コーパスに含まれている単語から辞書を作成します。
この際、出現頻度の低い単語に関してはという未知語を意味する記号にひとまとめにしてしまいます。
今回の実装では、出現回数が3回以下の単語は全てとして扱うようにしました。

私の実装は下記の通りです。
この関数は大まかに3つの部分に分かれており、それぞれ(1)コーパス中の単語の出現回数をカウントする部分、(2)出現回数の低い単語をにまとめる部分、(3)辞書をファイルに保存する部分で構成されています。

# コンストラクタで次の値が定義されている
    self.unknown_word_threshold = 3
    self.unknown_word_symbol = "<unk>"

def build_dict(self):
    # コーパス全体を見て、単語の出現回数をカウントする
    counter = collections.Counter()
    for filename in glob.glob(self.corpus_files):
        with open(filename, "r", encoding=self.corpus_encoding) as f:
            # Word breaking
            text = f.read()
            text = text.lower().replace("\n", " ")
            text = re.sub(r"[^a-z '\-]", "", text)
            text = re.sub(r"[ ]+", " ", text)

            # Preprocessing: Ignore a word starting with '-'
            words = [word for word in text.split() if not word.startswith("-")]
            counter.update(words)

    # 出現頻度の低い単語をひとつの記号にまとめる
    word_id = 0
    dictionary = {}
    for word, count in counter.items():
        if count <= self.unknown_word_threshold:
            continue

        dictionary[word] = word_id
        word_id += 1
    dictionary[self.unknown_word_symbol] = word_id

    print("# of unique words:", len(dictionary))

    # 辞書をpickleを使って保存しておく
    with open(self.dictionary_filename, "wb") as f:
        pickle.dump(dictionary, f)
        print("Dictionary is saved to", self.dictionary_filename)

    self.dictionary = dictionary

モデルの定義

モデル自体はも文字レベルの言語モデルを作った時と大差ありません。
以前はBasicRNNCellを使用していましたが、今回はBasicLSTMCellを使ってRNNのユニットをLSTMにしています。
学習させるときに入力する単語列の長さは、self.chunk_sizeで定義しており、今回は5としています。

def inference(self, input_data):
    """
    input_data: (batch_size, chunk_size, 1)
    initial_state: (batch_size, hidden_layer_size)
    """
    hidden_w = tf.Variable(tf.truncated_normal([self.input_layer_size, self.hidden_layer_size], stddev=0.01, dtype=tf.float32))
    hidden_b = tf.Variable(tf.ones([self.hidden_layer_size], dtype=tf.float32))
    output_w = tf.Variable(tf.truncated_normal([self.hidden_layer_size, self.output_layer_size], stddev=0.01, dtype=tf.float32))
    output_b = tf.Variable(tf.ones([self.output_layer_size], dtype=tf.float32))

    input_data = tf.one_hot(input_data, depth=self.input_layer_size, dtype=tf.float32) # (batch_size, chunk_size, input_layer_size)
    input_data = tf.reshape(input_data, [-1, self.chunk_size, self.input_layer_size])
    input_data = tf.transpose(input_data, [1, 0, 2]) # (chunk_size, batch_size, input_layer_size)
    input_data = tf.reshape(input_data, [-1, self.input_layer_size]) # (chunk_size * batch_size, input_layer_size)
    input_data = tf.matmul(input_data, hidden_w) + hidden_b # (chunk_size, batch_size, hidden_layer_size)
    input_data = tf.split(0, self.chunk_size, input_data) # chunk_size * (batch_size, hidden_layer_size)

    lstm = tf.nn.rnn_cell.BasicLSTMCell(self.hidden_layer_size, forget_bias=self.forget_bias)
    outputs, states = tf.nn.rnn(lstm, input_data, initial_state=lstm.zero_state(self.batch_size, tf.float32))

    # The last output is the model's output
    output = tf.matmul(outputs[-1], output_w) + output_b
    return output

def loss(self, model, labels):
    labels = tf.one_hot(labels, depth=self.output_layer_size, dtype=tf.float32)
    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(model, labels))
    return cost

def training(self, cost):
    optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(cost)
    return optimizer

学習を回す部分

モデルの定義が終わったので、次は学習を回す部分を定義します。
trainという関数が学習を回している部分になります。
ここの実装も前回と大きく変わってはいないかと思いますが、学習の進行度を見るためにイテレーションを回すごとにパープレキシティを計算して表示するようにしています。

また、実際にコーパスからデータを読み込んでバッチごとに区切ったデータを用意しているのはgenerate_batchという関数です。
イテレータとして実装しているので、generate_batch()に対してforループを回してあげると、コーパス全体のデータが順に得られるようになっています。

学習させてみる

以上で実装が終わったので、実際に学習させてみます。
先述の通り、コーパスはnewspaper_newswireというフォルダのテキストを使用します。

辞書の作成

まずは辞書を作成します。

# of unique words: 1146
Dictionary is saved to ./data/newspaper.dic

出現頻度が低いものを除くと、ユニークな単語は1146個存在したようです。

学習の実行

いよいよ学習を走らせます。意図的にコーパスを小さめにしているので、我が家のCPUでも100エポックが30分程度？で終わりました。
結果はこんな感じ。

Dictionary is successfully loaded
Dictionary size is: 1146
Epoch 1 completed out of 100; Loss = 4066.596182346344; Perplexity = 209.16334991551167
Epoch 2 completed out of 100; Loss = 3821.0914855003357; Perplexity = 145.92456359526972
Epoch 3 completed out of 100; Loss = 3641.759223461151; Perplexity = 114.590039718061
...(中略)...
Epoch 100 completed out of 100; Loss = 1975.9882043004036; Perplexity = 13.141435590739214
Trained model is saved to ./data/newspaper.ckpt

パープレキシティも13程度まで下がり、学習がうまく進んでいる様子が確認できます。

問題点

今回はTensorFlowのBasicLSTMCellというセルを使用して実装しました。
この実装の問題点は、self.chunk_sizeで定義した長さの単語列を与えないとモデルからの出力が得られない点です。（私の知る限りは、ですが。。もし何か手段があったらご教示ください。。）
可変長の単語列を与えようと思うと、なにやらtf.nn.dynamic_rnnなるものを使わないといけないようなのですが、これに関してはまだ詳しく調査できていません。
参考：Variable Sequence Lengths in TensorFlow

今回と同じモデルをCNTKを使っても実装してみたのですが、そちらの場合は可変長の単語列も入力として使えました。
一度CNTKの方の実装を紹介してから、余力があったらdynamic_rnnの使い方も調べてみようと思っています。

今日も窓辺でプログラム

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