今日も窓辺でプログラム

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

Tensor2Tensorを使って独自データでseq2seqしてみる

はじめに

Tensor2Tensorという、広く使われているモデルやデータセットが実装されているディープラーニングのライブラリがあります。
GitHubのREADMEによると、Tensor2TensorはGoogle Brainチームなどによって開発されているようで、中身はTensorFlowで実装されています。

用意されているモデルとデータセットを使うだけなら、コーディングなしでディープラーニングを試すことができます。
モデルやデータセットの追加も容易に行えるようになっていて感動したのですが日本語の記事はあまり見当たらなかったので、せっかくなので今回試した「自分で用意したデータでseq2seqを学習させる」ための使い方を紹介しようと思います。

環境はUbuntu 16.04のでGPUを使用した学習を行いました。

何ができるの?

Tensor2Tensorは、自分で用意したデータを読みこむ部分を実装するだけで、主要モデルでのディープラーニングが簡単に行えます。
今回私はseq2seqで機械翻訳や対話モデルの作成を行ったのですが、単語分割もwordpieceを使って自動的に面倒を見てくれるので、MeCab等を使用して分かち書きしておく、といった作業も必要ありません。必要なのは、入力と出力のペア、それだけです。

wordpieceについては、以前当ブログでも試してみたことがあるので、興味のある方はそちらの記事もどうぞ。
SentencePieceを使用してRNN言語モデルを学習させてみる - 今日も窓辺でプログラム

インストール

Tensor2Tensorを使うには、まずTensorFlowをインストールしないといけません。公式のドキュメントに沿ってインストールを進めてください。
学習にGPUを使いたい場合は、当然GPU版のTensorFlowをインストールしておく必要があります。
Installing TensorFlow  |  TensorFlow

TensorFlowのインストールが終わったら、Tensor2Tensorをpipからインストールします。

pip install tensor2tensor

まずは用意されたモデルとデータセットでseq2seqしてみる

Walkthroughをそのまま試す

Walkthroughに書いてある例が、ちょうど英語からドイツ語への機械翻訳を学習・評価する例ですので、まずはそのままなぞってみました。省略しているコマンドも多数あるので、実際に実行する場合は元のWalkthroughも参照しながら行てください。

まずは解くべき問題(=データセット)や、使用するモデル、ハイパーパラメータを定義します。

PROBLEM=translate_ende_wmt32k
MODEL=transformer
HPARAMS=transformer_base_single_gpu

translate_ende_wmt32kというのが機械翻訳用のデータセット、transformerというのがAttention is all you needというGoogleの論文で発表されたTransformerというモデルです

次に学習に使用するデータを生成します。下記のt2t-datagenコマンドを実行すると、$PROBLEMを参照してデータセットの生成(必要ならローカルファイルやネットワーク越しに読み込む)し、wordpieceによるトークナイズまで行ってくれます。その結果は$DATA_DIRに保存されます。

t2t-datagen \
  --data_dir=$DATA_DIR \
  --tmp_dir=$TMP_DIR \
  --problem=$PROBLEM

次は学習です。先ほどの引数に加え、ネットワークのモデルを$MODELで、そのハイパーパラメータを$HPARAMSで与えています。学習結果は$TRAIN_DIRに出力されます。
私の環境だと、この学習に半日程度かかった記憶があります。

t2t-trainer \
  --data_dir=$DATA_DIR \
  --problem=$PROBLEM \
  --model=$MODEL \
  --hparams_set=$HPARAMS \
  --output_dir=$TRAIN_DIR

ここまでで学習は終わりです。あとはWalkthroughにある形でファイルからのデコードを行ってもよいですし、BLEUスコアを計算してもよいです。
私の場合はインタラクティブにモデルを試してみたかったので、下記のコマンドを入力しました。

BEAM_SIZE=4
ALPHA=0.6

t2t-decoder \
   --data_dir=$DATA_DIR \
   --problem=$PROBLEM \
   --model=$MODEL \
   --hparams_set=$HPARAMS \
   --output_dir=$TRAIN_DIR \
   --decode_hparams="beam_size=$BEAM_SIZE,alpha=$ALPHA" \
   --decode_interactive=true

しばらく待つと学習済みモデルがロードされ、対話的にモデルを試せるようになります。

モデルを変えてみる

実装されているモデルであれば、モデルを変えるのは非常に簡単です。$MODELと$HPARAMSを下記のように使用したいモデルのものに変更して、もう一度t2t-datagenコマンドから順に走らせてあげるだけです。下記パラメータでは、LSTMベースのモデルが試せます。

MODEL=lstm_seq2seq_attention_bidirectional_encoder
HPARAMS=lstm_luong_attention_multi

LSTMでもいくつものモデルが実装されてます。登録されているモデルやハイパーパラメータの一覧は下記コマンドで確認できます。日々増え続けているようです。

t2t-trainer --registry_help

自分で用意したデータで学習してみる

学習させたいデータセット

今回使用するデータとして、下記のようなイメージの2カラムのTSVファイルを事前に用意していました。最初の行はヘッダーです。
データはお見せできないですが、日本語の対話データを集めたコーパスです。

question	answer
こんにちは	こんにちは!
調子どう?	今日は絶好調です。
...

Problemの定義

Tensor2Tensorを使った学習や評価は、上記のようにコマンドで行います。使用するデータセットを変えたい場合は、$PROBLEMの値を英語ドイツ語翻訳のものから、自分で用意したデータセットを指すようにしてあげる必要があります。
方法は下記ドキュメントに書かれている通り、Problemというクラスを継承したクラスを実装してあげる必要があります。
T2T: Train on Your Own Data | tensor2tensor

私の用意したTSVの場合、用意すべきスクリプトは下記のようなものになります。
Textでのseq2seqの場合は、Text2TextProblemというクラスを継承すると簡単にデータセットを定義できます。

import re
import pandas as pd

from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry

@registry.register_problem
class MyProblem(text_problems.Text2TextProblem):
  """Predict next line of poetry from the last line. From Gutenberg texts."""

  @property
  def approx_vocab_size(self):
    return 2**13  # ~8k

  @property
  def is_generate_per_split(self):
    # generate_data will shard the data into TRAIN and EVAL for us.
    return False

  @property
  def dataset_splits(self):
    """Splits of data to produce and number of output shards for each."""
    # 10% evaluation data
    return [{
        "split": problem.DatasetSplit.TRAIN,
        "shards": 9,
    }, {
        "split": problem.DatasetSplit.EVAL,
        "shards": 1,
    }]

  def generate_samples(self, data_dir, tmp_dir, dataset_split):
    filename = '<TSV path>'
    df = pd.read_csv(filename, delimiter='\t')
    for key, row in df.iterrows():
        q, a = row['question'], row['answer']
        if not q or not a:
            continue

        yield {
            'inputs': q.strip(),
            'targets': a.strip()
        }

ドキュメントのサンプルから手を入れているのは、generate_samplesの関数の中身くらいです。TSVから入力と出力のペアを読み込み、yieldで順に返しているだけです。10行程度書いただけでしょうか。
このファイルをmy_problem.pyとして保存しておきます。
また、同一フォルダに下記の内容の__init__.pyも保存しておく必要があります。

from . import my_problem

学習してみる

先ほど作成した2つのPythonスクリプトを保存したフォルダを$USR_DIR変数に設定した上で、下記コマンドを実行すると学習は完了です。どちらのコマンドにも、--t2t_usr_dirというオプションが増えており、これは自分で用意したスクリプトをコマンドに見えるようにするためのものです。

t2t-datagen \
  --data_dir=$DATA_DIR \
  --tmp_dir=$TMP_DIR \
  --problem=$PROBLEM \
  --t2t_usr_dir=$USR_DIR

t2t-trainer \
  --data_dir=$DATA_DIR \
  --problem=$PROBLEM \
  --model=$MODEL \
  --hparams_set=$HPARAMS \
  --output_dir=$TRAIN_DIR \
  --t2t_usr_dir=$USR_DIR

評価の場合も同様に--t2t_usr_dirオプションを付与しておけばよく、例えば下記のコマンドでは学習したモデルを対話的に試すことができます。

t2t-decoder \
   --data_dir=$DATA_DIR \
   --problem=$PROBLEM \
   --model=$MODEL \
   --hparams_set=$HPARAMS \
   --output_dir=$TRAIN_DIR \
   --decode_hparams="beam_size=$BEAM_SIZE,alpha=$ALPHA" \
   --decode_interactive=true \
   --t2t_usr_dir=$USR_DIR

こちらも結果はお見せできないのですが、割と良い感じに対話できるモデルが出来上がりました。


おわりに

今回は紹介しませんでしたが、TensorFlowのServerを使って本番運用の環境をすぐ整えられたりもするようです。
実験レベルで様々のモデルやパラメータを試すのには非常に便利ですし、本番運用まで持っていけそうなのでTensor2Tensorは魅力的なオプションなのではないでしょうか。