今日も窓辺でプログラム

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

RNN言語モデルのpytorch実装をWikipediaの記事で学習させてみる

はじめに

pytorchのGitHubに上がっているサンプルを見ていたら、RNNを使って言語モデルを実装しているものがありました。
examples/word_language_model at master · pytorch/examples · GitHub

本当はいろんなネットワークを1から実装するのがよいのでしょうが、そのような経験もあまりないのでまずはサンプルを写経して走らせてみようと思いました。
最近Wikipediaのデータを単語分割したコーパスを用意したので、そのコーパスに対して上記実装を走らせてみます。

Jupyter Notebook

今回の実験で使用したJupyter NotebookをGitHubにアップロードしてあります。ほとんど公式のサンプルそのままですが、記事を読む際の参考にしていただければと思います。
github.com

下準備

まずはコーパスとなるWikipediaの記事を準備します。基本的には以前書いた記事で使用したコードを使用します。
1点だけ、分かち書きされた文章を出力する際に、文の始まりと終わりにそれぞれを追加するような変更を入れてあります(変更箇所はここ)。

このコードを走らせると、wakati_corpus.txtという名前で

<bos> 悔返 <eos>
<bos> 悔返 ( くい かえし ) と は 、 中世 日本 において 、 和 与 ・ 寄進 など の 財産 処分 を 行っ て 所有 権 の 移動  が 行わ れ た 後 に 元 の 所有 者 あるいは その 子孫 ら が その 行為 を 否定 し て 取り戻す 行為 。 <eos>
...

というようなファイルが保存されます。前回の記事でも書きましたが、記事を文に分割するのが面倒だったため、記事中の1行を1つの文章として扱っています。

データ読み込み部分

サンプルコードには、コーパスに出てくる単語一覧を保持するDictionaryクラスと、コーパスとなるテキストを単語IDの列として保持するCorpusクラスが定義されています。

Dictionaryクラスはサンプルのものがそのまま使用できますが、CorpusクラスはWikipediaの記事を読み込むように変えてあげないといけません。
今回は以下のように実装してみました。corpus_path (= wakati_corpus.txt) から10,000行読み込んで*1、8:1:1の割合でトレーニング:検証:テストデータに分割しました。

# wakati_corpus.txtから一部データを読み込む
class Corpus(object):
    def __init__(self, corpus_path):
        self.dictionary = Dictionary()

        texts = self.load(corpus_path, 10000)
        texts = self.tokenize(texts)
           
        train_size = int(len(texts) * 0.8)
        valid_size = int(len(texts) * 0.9)
        
        self.train = texts[:train_size]             # 8割がトレーニングデータ
        self.valid = texts[train_size+1:valid_size] # 1割が検証用データ
        self.test = texts[valid_size+1:]            # 1割がテストデータ
        
        self.train = self.assign_ids(self.train)
        self.valid = self.assign_ids(self.valid)
        self.test = self.assign_ids(self.test)

    def load(self, corpus_path, count):
        texts = []
        with open(corpus_path, 'r', encoding='utf-8') as f:
            i = 0
            line = f.readline()
            while line:
                line = line.strip()
                texts.append(line)
                line = f.readline()
                i += 1
                if i > count: break
        return texts

    def tokenize(self, texts):
        tokens = []
        for text in texts:
            text = text.lower()
            words = text.split()

            for word in words:
                self.dictionary.add_word(word)
                
            tokens.extend(words)
            
        return tokens
    
    def assign_ids(self, texts):
        tokens = len(texts)
        ids = torch.LongTensor(tokens)
        token = 0
        for word in texts:
            ids[token] = self.dictionary.word2idx[word]
            token += 1
            
        return ids

このクラスが初期化された時点で、self.train, self.valid, self.testはそれぞれ単語IDの列をtorch.LongTensorとして保持している状態になります。

今回サンプルコードと違うのはこの読み込み部分くらいで、実は残りはほぼサンプルそのままです。。

モデルの定義

RNNのモデルはRNNModelというクラスで定義されているので、そのまま使用します。
語彙数次元の入力をEmbedding層で密なテンソルにし、任意の数の隠れ層を経た後、語彙数次元の出力層につながっています。

Dropoutの実装や重み・隠れ層の初期化の実装など、非常に参考になります。
各層の次元数や隠れ層の層数、LSTMを使うかGRUを使うかなどがインスタンス化するときに選択できる実装になっていて便利です。

この関数を使い、今回は次のような設定で試してみることにしました。Embeddingの次元数は200、隠れ層は200次元1層で、GRUユニットを使用します。コスト関数はクロスエントロピーです。

emsize=200
nhid=200
nlayers=1
dropout=0.5
tied=False
ntokens = len(corpus.dictionary)
model = RNNModel('GRU', ntokens, emsize, nhid, nlayers, dropout, tied)
if use_gpu:
    model.cuda()

criterion = nn.CrossEntropyLoss()

バッチ化の準備

Corpusクラスで読み込んだ単語ID列をバッチ処理するためにbatchify関数を使い次のような変換をかけるようです。
テキストが a, b, c, d, ..., x, y, zというテキストで、バッチサイズが4の場合、batchify関数の出力は
 \displaystyle \left[ \begin{array} {rrrr}
a & g & m & s \\
b & h & n & t \\
c & i & o & u\\
d & j & p & v \\
e & k & q & w \\
f & l & r & x \\
\end{array} \right]
となります。(本当は各要素は文字ではなく単語ID)
実際に学習のイテレーションを回すときに、各行から単語ID列を取り出すことで、バッチ化された入力が用意できるわけですね。

この関数を使い、バッチサイズをどの程度にすればよいかのノウハウもないので、ひとまずサンプルのデフォルト値である学習時は20、検証・テスト時は10というサイズに設定しました

# コーパスを読み込む
corpus = Corpus('./wakati_corpus.txt')

# 各データをバッチに分割
batch_size = 20
eval_batch_size = 10
train_data = batchify(corpus.train, batch_size)
val_data = batchify(corpus.valid, eval_batch_size)
test_data = batchify(corpus.test, eval_batch_size)

学習部分のコード

1エポック分の学習をするコードがtrain関数で定義されています。
bathify関数で分割した単語ID列から入力となるシーケンスを取り出す部分はget_batch関数に切り分けられています。

また、誤差の逆伝播を行うときに、バッチ1個分のシーケンス分だけ逆伝播させるためにrepackage_hidden関数を使ってそれまでのヒストリーと切り離すという作業も行われています。新しいVariableを使わずに同じVariableを使いまわしてしまうと、常にトレーニングデータの一番最初の部分まで逆伝播をしてしまうようです。

学習時にclip_grad_normという関数を呼び出している箇所があります。
これはRNNやLSTMの勾配爆発を抑制するための対策で、Gradient clippingと呼ばれているようです。私も詳しく理解したわけではないのですが、目的関数の勾配が急になると、パラメータの更新時にパラメータが大きく動いてしまいます。Gradient clipplingでは、勾配の大きさが閾値を超えたとき、方向は変えずに大きさを閾値に収まるようにスケールすることで勾配が大きくなりすぎないようにしているようです。
EnVision: Deep Learning : Why you should use gradient clipping

イテレーションを回す

イテレーションを回す部分もサンプルをそのまま拝借しました。
今回は、逆伝播させるシーケンスの長さを35にして、ひとまず30エポック回す設定で学習を走らせてみました。

すると1エポック目が↓のような状態から始まり、

| epoch   1 |   200/  626 batches | lr 20.00 | ms/batch 63.18 | loss  7.47 | ppl  1752.46
| epoch   1 |   400/  626 batches | lr 20.00 | ms/batch 63.78 | loss  6.26 | ppl   524.86
| epoch   1 |   600/  626 batches | lr 20.00 | ms/batch 62.81 | loss  5.98 | ppl   395.01
-----------------------------------------------------------------------------------------
| end of epoch   1 | time: 41.42s | valid loss  5.99 | valid ppl   397.75
-----------------------------------------------------------------------------------------

20エポックを終えたころには↓ぐらいまでロスが小さくなっていました。

| epoch  20 |   200/  626 batches | lr 0.08 | ms/batch 70.54 | loss  4.23 | ppl    68.94
| epoch  20 |   400/  626 batches | lr 0.08 | ms/batch 67.30 | loss  4.09 | ppl    59.60
| epoch  20 |   600/  626 batches | lr 0.08 | ms/batch 64.89 | loss  3.97 | ppl    53.00
-----------------------------------------------------------------------------------------
| end of epoch  20 | time: 43.87s | valid loss  5.56 | valid ppl   259.54
-----------------------------------------------------------------------------------------

この後も少し学習を進めてみたのですが、学習率もどんどん小さくなっていき、ロスにもほとんど変化がなくなったので、25エポックあたりで打ち切りました。

テストデータでの評価

評価用の関数も用意されているので使わせてもらいましょう。
このようにmodel.eval()を実行すると、評価用モードになりdropoutが無効になるようです。

with open('model.pt', 'rb') as f:
    model = torch.load(f)

test_loss = evaluate(test_data)
print('=' * 89)
print('| End of training | test loss {:5.2f} | test ppl {:8.2f}'.format(
    test_loss, math.exp(test_loss)))
print('=' * 89)

テストデータでも、トレーニング・検証用と同水準(というか一番良い)のロスに落ち着いていることが確認できました。

=========================================================================================
| End of training | test loss  5.29 | test ppl   198.01
=========================================================================================

文章を生成してみる

学習した言語モデルをもとに文章を生成するサンプルコードも公開されているので、それをもとに文章を生成してみます。

隠れ層の状態はRNNModel.init_hidden()を使用してランダムに初期化、を最初の入力にしてが出てくるまで次の単語を予測し続ける、ということを行います。
モデルの出力は一番確率の高い単語を採用するのではなく、torch.multinomial関数を使い、出力の指数を取った値で定義される多項分布からサンプリングを行うことで次に来る単語を予測しているようです。(こういう方法は一般的なんでしょうか??)

# ランダムな文章を生成してみる
temperature=1.0
for j in range(10):
    # 隠れ層はランダムに初期化
    hidden = model.init_hidden(1)
    
    # <bos>から始める
    input = Variable(torch.LongTensor([[corpus.dictionary.word2idx['<bos>']]]), volatile=True).cuda()
    model = model
    results = []
    for i in range(100):
        output, hidden = model(input, hidden)
        word_weights = output.squeeze().data.div(temperature).exp()
        word_idx = torch.multinomial(word_weights, 1)[0]
        input.data.fill_(word_idx)
        word = corpus.dictionary.idx2word[word_idx]

        # <eos>が出たら1文終わり
        if word == '<eos>':
            break

        results.append(word)

    print(' '.join(results))

上記コードで10個の文章を生成してみた結果の一例がこちらです。

1 縄 の きのこ と なる 。
天平 15 年 ( で 1575 年 ) 証さ 三恵子 、 展開 さ れ た コンチネンタル ・ 吉浦 低湿 で 、 レニングラード 質 が 盛ん に 07 から 埋め込む こと を 決め た 。 高等 学校 に 記す こと も ある 。 しかし 、 1930 年代 から 復活 し て 逝去 という 事情 は 254 年 の 各 アナウンサー が 故障 を 行い 、 1973 年 ( 昭和 24 年 ) に は 69 歳 で くれ 荘 の 大 規模 だっ た と さ れる と 、 政府 は 新た に あろ う 、 選集 布 ラッヘンマン の 高木
学術 対策 が 1 回 裏 の 規定 により 得 られ た 。
難し さ も ティナ・ターナー の 各所 として オランダ ウェルター 組 の 2 回戦 を ディーゼル 機関 車 とともに 、 530 女 に は パッド を ぶつけ 守り抜い ながら 、 補強 さ れる 2000 年 も レ 天皇 にて 古谷 二 冠 と する ユダヤ 人 を 利用 する 。
本 は ガス を 充分 な 4 倍 で 、 アンソロジスト で 行わ れる 。 ここ で も その よう な 操縦 を 下さ れ て いる 。 4 回 限定 シングル 公式 戦 が 同所 し た 。
コスト に ステイシー・トーテン
烏 で は 距離 、 索道 と 同様 、 とても 説明 を 吸収 ず 、 末弟 も 多く の 学齢 甥 を 支え て いる 。
俄然 2 人 を うがち の 男女 も メッセージ を 元帥 する こと が あっ た 。
揃える と 、 無名 な cd に 乗っ て おもちゃ へ の キャスティング ・ 美貌 の どんな プレイオフ を もつ と さ れる 。
浦和 の 一部 に 所在 する ふり し た 。

うーん。日本語っぽいところもありますが、全然なところがほとんどですね。また、元の文章が長いからか、なかなかに到達せず長い文章が生成されがちなようでした。



*1:全記事分となると時間もかかりますし、長時間学習回し続けられる環境持っていないのが理由です。。