fairseqを使用して今まで消費していた時間を取り戻す

（これまで失った時間を取り戻せるはず）ないです。

この記事はマイナビ Advent Calendar 2021 の8日目となります。

以前、弊社のアドカレでチャットボットについて書きました。今回もそれに習いfairseqについて書きたいと思います。

はじめに

以前作成したチャットボットのチューニングに勤しみつつ、そろそろTransformerとか試してみたいなぁと思い、早2年ぐらい立ちました。

いつものようにTwitterで情報収集を努めていたところ、以下のツイートを発見。

去年の対話システムライブコンペで大いなる力を見せつけて優勝した、NTTさんの16億パラメータ日本語対話モデルがついに公開。待ってました・・・！🎉🎉🎉
日本語対話システム研究がこれからどんどん活発化しそうですね。https://t.co/tVXGE7H56Hhttps://t.co/gVA5JkKcSE

— Seitaro Shinagawa (@sei_shinagawa) September 20, 2021

「すごい、ちょうど欲しかったやつ！Githubでも公開されているし、モデルの参考にしよう！」と思ってGithubにアクセスしたところ、何やら見慣れないコマンドラインツールがありました。

それがfairseqとのファーストコンタクトとなります。

fairseq

fairseqとはFacebook AI Research（FAIR）が出しているPyTorch向けのシーケンスモデル用ツールキットです。翻訳や要約、言語モデル、テキスト生成タスクなどで利用するモデルの訓練や推論を高速にイテレーションできるよう簡単化するためのツールとなります。

マルチGPUによる分散トレーニングや高速なビームサーチなど様々なオプションが提供されており、正直なところオプション数が多すぎて把握しきれていません。もう少し公式ドキュメントを読み込もうかと思います。

導入は以下の3行です。簡単ですね。

$ git clone https://github.com/pytorch/fairseq
$ cd fairseq
$ pip install --editable ./

今回、こちらのfairseqを使用して、これまで苦楽をともにしてきたSeq2Seqの学習を試してみたいと思います（以下のチュートリアルを参考に進めていきます）。

実際に使ってみる

モデルについて

fairseq自体はPyTorchを拡張している形で実装されているため、これまでPyTorchで実装していたモデルをほぼそのまま流用できます。この時点で好きになっちゃいそうです。

今回は簡単な1層GRUによるEncoder-Decoderモデルで学習をやってみたいと思います。

• 1層GRUによるEncoder

./models/gru.py

import torch.nn as nn
from fairseq import utils
from fairseq.models import FairseqEncoder


class SimpleGRUEncoder(FairseqEncoder):

    def __init__(
            self, args, dictionary, embed_dim=128, hidden_dim=128, nlayers=1, dropout=0.5, w=None):
        super().__init__(dictionary)

        self.args = args

        # encoder embedding
        self.encoder_embed = nn.Embedding(
            len(dictionary), embed_dim, padding_idx=dictionary.pad())

        # 事前学習済み言語モデルがある場合はWを固定
        if w is not None:
            self.encoder_embed.weight.data.copy_(w)
            self.encoder_embed.requireds_grad_ = False

        # encoder
        self.encoder = nn.GRU(
            embed_dim,  # 入力次元
            hidden_dim,  # 隠れ層
            batch_first=True,
            dropout=dropout,
            num_layers=nlayers,
            bidirectional=False
        )
        # dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self, src_tokens, src_lengths):

        if self.args.left_pad_source:
            src_tokens = utils.convert_padding_direction(
                src_tokens,
                padding_idx=self.dictionary.pad(),
                left_to_right=True
            )

        # embedding
        embed_src = self.dropout(self.encoder_embed(src_tokens))  # encode embedding
        embed_src = nn.utils.rnn.pack_padded_sequence(embed_src, src_lengths, batch_first=True)

        # encoder
        _output, final_hidden = self.encoder(embed_src)  # encoder output

        return {
            'final_hidden': final_hidden.squeeze(0),
        }

    def reorder_encoder_out(self, encoder_out, new_order):
        final_hidden = encoder_out['final_hidden']
        return {
            'final_hidden': final_hidden.index_select(0, new_order)
        }

• 1層GRUと出力層からなるDecoder

gru.py

import torch.nn as nn
from fairseq import utils
from fairseq.models import FairseqEncoder


class SimpleGRUDecoder(FairseqDecoder):

    def __init__(self, dictionary, embed_dim, encoder_hidden=128, decoder_hidden=128, n_layer=1, dropout=0.5, w=None):
        super().__init__(dictionary)

        # decoder embedding
        self.decoder_embed = nn.Embedding(
            len(dictionary), embed_dim, padding_idx=dictionary.pad()
        )

        # 事前学習済み言語モデルがある場合はWを固定
        if w is not None:
            self.decoder_embed.weight.data.copy_(w)
            self.decoder_embed.requireds_grad_ = False

        # decoder
        self.decoder = nn.GRU(
            input_size=encoder_hidden + decoder_hidden,
            hidden_size=decoder_hidden,  # 隠れ層
            dropout=dropout,
            num_layers=n_layer,
            bidirectional=False
        )

        # dropout
        self.dropout = nn.Dropout(dropout)

        # 出力層
        self.output = nn.Linear(decoder_hidden, len(dictionary))

    def forward(self, prev_output_tokens, encoder_out):

        bsz, tgt_len = prev_output_tokens.size()
        final_encoder_hidden = encoder_out['final_hidden']

        # embedding
        embed_trg = self.dropout(self.decoder_embed(prev_output_tokens))  # decoder embedding
        decoder_input = torch.cat(
            [embed_trg, final_encoder_hidden.unsqueeze(1).expand(bsz, tgt_len, -1)], dim=2)

        initial_state = final_encoder_hidden.unsqueeze(0)

        # decoder
        dec_output, _ = self.decoder(
            decoder_input.transpose(0, 1), 
            initial_state)  # decoder output

        decoder_input = dec_output.transpose(0, 1)
        decoder_input = self.output(decoder_input)

        return decoder_input, None

PyTorchとの違いとしては
* FairseqEncoderとFairseqDecoderのinitにdictionaryを渡している
* FairseqEncoderでreorder_encoder_outを実装しなければならない

ぐらいでしょうか。fairseqの内部までまだ読んでいないので、モデルクラスに語彙データを与える理由については理解できていません。後の工程で使用するコマンドたちでデータをそのまま渡しているので、fairseq内部で自動的に辞書化した上でモデルに渡されているかもしれません。

もしそうだとしたらなんと便利なことでしょうか。

コマンドラインツールへモデルを登録し学習してみる

モデルを定義したらいよいよコマンドラインツールで実行できるようにしていきます。

gru.py

@register_model('simple_gru')
class SimpleGRUModel(FairseqEncoderDecoderModel):

    @staticmethod
    def add_args(parser):
        """add_argsでコマンドラインオプションを登録"""

        parser.add_argument(
            '--encoder-embed-dim', type=int, metavar='N',
            help='dimensionality of the encoder embeddings',
        )
        parser.add_argument(
            '--encoder-hidden-dim', type=int, metavar='N',
            help='dimensionality of the encoder hidden state',
        )
        parser.add_argument(
            '--encoder-dropout', type=float, default=0.1,
            help='encoder dropout probability',
        )
        parser.add_argument(
            '--decoder-embed-dim', type=int, metavar='N',
            help='dimensionality of the decoder embeddings',
        )
        parser.add_argument(
            '--decoder-hidden-dim', type=int, metavar='N',
            help='dimensionality of the decoder hidden state',
        )
        parser.add_argument(
            '--decoder-dropout', type=float, default=0.1,
            help='decoder dropout probability',
        )

    @classmethod
    def build_model(cls, args, task):
        """コマンドラインオプションで渡された値を元に学習モデルを初期化"""

        # EncoderとDecoderの初期化
        encoder = SimpleGRUEncoder(
            args=args,
            dictionary=task.source_dictionary,
            embed_dim=args.encoder_embed_dim,
            hidden_dim=args.encoder_hidden_dim,
            dropout=args.encoder_dropout,
        )
        decoder = SimpleGRUDecoder(
            dictionary=task.target_dictionary,
            encoder_hidden=args.encoder_hidden_dim,
            embed_dim=args.decoder_embed_dim,
            decoder_hidden=args.decoder_hidden_dim,
            dropout=args.decoder_dropout,
        )
        # EncoderとDecoderを持つSeq2Seqを定義
        model = SimpleGRUModel(encoder, decoder)

        # モデルアーキテクチャを出力してみる
        print(model)

        return model


@register_model_architecture('simple_gru', 'tutorial_simple_gru')
def tutorial_simple_gru(args):
    """--archで指定できるようにモデルアーキテクチャを登録"""
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 256)
    args.encoder_hidden_dim = getattr(args, 'encoder_hidden_dim', 256)
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 256)
    args.decoder_hidden_dim = getattr(args, 'decoder_hidden_dim', 256)

すごい！簡単！
fairseq-train自体のオプションの他に独自にオプションを登録できるあたりが有能だと思います。私は好きです。

こうして出来上がったgru.pyをモデル学習用のfairseq-trainコマンドで実行してみます。

$ fairseq-train ../data-bin/iwslt14.tokenized.de-en \
  --arch tutorial_simple_gru \
  --encoder-dropout 0.2 --decoder-dropout 0.2 \
  --optimizer adam --lr 0.005 --lr-shrink 0.5 \
  --max-tokens 12000 \
  --max-epoch 1

...
fairseq-train: error: argument --arch/-a: invalid choice: 'tutorial_simple_gru' (choose from 's2t_berard', 's2t_berard_256_3_3', 's2t_berard_512_3_2', 

--archで指定したtutorial_simple_gruは無効な選択だよ、というエラーが発生しました。こういうときは焦らずググりましょう。

どうやらfairseq/fairseq/models/*に実装したスクリプトを配置しなければならないそうです。移動して再実行してみましょう。

 $ fairseq-train ../fairseq/data-bin/iwslt14.tokenized.de-en \
  --arch tutorial_simple_gru \
  --encoder-dropout 0.2 \
  --decoder-dropout 0.2 \
  --optimizer adam --lr 0.005 --lr-shrink 0.5 \
  --max-tokens 12000 \
  --max-epoch 1
2021-12-06 00:30:01 | INFO | fairseq.tasks.text_to_speech | Please install tensorboardX: pip install tensorboardX
2021-12-06 00:30:03 | INFO | fairseq_cli.train | {'_name': None, 'common': {'_name': None, 'no_progress_bar': False, 'log_interval': 100, 'log_format': None, 'log_file': None, 'tensorboard_logdir': None,
...
2021-12-06 00:30:03 | INFO | fairseq.tasks.translation | [de] dictionary: 8848 types
2021-12-06 00:30:03 | INFO | fairseq.tasks.translation | [en] dictionary: 6632 types
...
  "num_layers={}".format(dropout, num_layers))
SimpleGRUModel(
  (encoder): SimpleGRUEncoder(
    (encoder_embed): Embedding(8848, 256, padding_idx=1)
    (encoder): GRU(256, 256, batch_first=True, dropout=0.2)
    (dropout): Dropout(p=0.2, inplace=False)
  )
  (decoder): SimpleGRUDecoder(
    (decoder_embed): Embedding(6632, 256, padding_idx=1)
    (decoder): GRU(512, 256, dropout=0.2)
    (dropout): Dropout(p=0.2, inplace=False)
    (output): Linear(in_features=256, out_features=6632, bias=True)
  )
)
2021-12-06 00:30:03 | INFO | fairseq_cli.train | SimpleGRUModel(
  (encoder): SimpleGRUEncoder(
    (encoder_embed): Embedding(8848, 256, padding_idx=1)
    (encoder): GRU(256, 256, batch_first=True, dropout=0.2)
    (dropout): Dropout(p=0.2, inplace=False)
  )
  (decoder): SimpleGRUDecoder(
    (decoder_embed): Embedding(6632, 256, padding_idx=1)
    (decoder): GRU(512, 256, dropout=0.2)
    (dropout): Dropout(p=0.2, inplace=False)
    (output): Linear(in_features=256, out_features=6632, bias=True)
  )
)
2021-12-06 00:30:03 | INFO | fairseq_cli.train | task: TranslationTask
2021-12-06 00:30:03 | INFO | fairseq_cli.train | model: SimpleGRUModel
2021-12-06 00:30:03 | INFO | fairseq_cli.train | criterion: CrossEntropyCriterion
2021-12-06 00:30:03 | INFO | fairseq_cli.train | num. shared model params: 6,653,416 (num. trained: 6,653,416)
2021-12-06 00:30:03 | INFO | fairseq_cli.train | num. expert model params: 0 (num. trained: 0)
2021-12-06 00:30:03 | INFO | fairseq.data.data_utils | loaded 7,283 examples from: ../fairseq/data-bin/iwslt14.tokenized.de-en/valid.de-en.de
2021-12-06 00:30:03 | INFO | fairseq.data.data_utils | loaded 7,283 examples from: ../fairseq/data-bin/iwslt14.tokenized.de-en/valid.de-en.en
2021-12-06 00:30:03 | INFO | fairseq.tasks.translation | ../fairseq/data-bin/iwslt14.tokenized.de-en valid de-en 7283 examples
2021-12-06 00:30:03 | INFO | fairseq_cli.train | training on 1 devices (GPUs/TPUs)
2021-12-06 00:30:03 | INFO | fairseq_cli.train | max tokens per device = 12000 and max sentences per device = None
2021-12-06 00:30:03 | INFO | fairseq.trainer | Preparing to load checkpoint checkpoints/checkpoint_last.pt
2021-12-06 00:30:03 | INFO | fairseq.trainer | No existing checkpoint found checkpoints/checkpoint_last.pt
2021-12-06 00:30:03 | INFO | fairseq.trainer | loading train data for epoch 1
2021-12-06 00:30:03 | INFO | fairseq.data.data_utils | loaded 160,239 examples from: ../fairseq/data-bin/iwslt14.tokenized.de-en/train.de-en.de
2021-12-06 00:30:03 | INFO | fairseq.data.data_utils | loaded 160,239 examples from: ../fairseq/data-bin/iwslt14.tokenized.de-en/train.de-en.en
2021-12-06 00:30:03 | INFO | fairseq.tasks.translation | ../fairseq/data-bin/iwslt14.tokenized.de-en train de-en 160239 examples
2021-12-06 00:30:03 | INFO | fairseq.data.iterators | grouped total_num_itrs = 388
epoch 001:   0%|                                                                                                                                  | 0/388 [00:00<?, ?it/s]2021-12-06 00:30:03 | INFO | fairseq.trainer | begin training epoch 1
2021-12-06 00:30:03 | INFO | fairseq_cli.train | Start iterating over samples
/Users/***/lib/python3.6/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling
  warnings.warn('User provided device_type of \'cuda\', but CUDA is not available. Disabling')
epoch 001: 100%|▉| 387/388 [20:31<00:02,  2.83s/it, loss=6.677, ppl=102.33, wps=3156.9, ups=0.31, wpb=10135.2, bsz=398.5, num_updates=300, lr=0.005, gnorm=0.369, train_wa2021-12-06 00:50:38 | INFO | fairseq_cli.train | begin validation on "valid" subset

[W ParallelNative.cpp:214] Warning: Cannot set number of intraop threads after parallel work has started or after set_num_threads call when using native parallel backend (function set_num_threads)

2021-12-06 00:50:58 | INFO | valid | epoch 001 | valid on 'valid' subset | loss 6.286 | ppl 78.05 | wps 8848.2 | wpb 7766.2 | bsz 316.7 | num_updates 388                 
2021-12-06 00:50:58 | INFO | fairseq.checkpoint_utils | Preparing to save checkpoint for epoch 1 @ 388 updates
2021-12-06 00:50:58 | INFO | fairseq.trainer | Saving checkpoint to checkpoints/checkpoint1.pt
2021-12-06 00:50:58 | INFO | fairseq.trainer | Finished saving checkpoint to checkpoints/checkpoint1.pt
2021-12-06 00:50:59 | INFO | fairseq.checkpoint_utils | Saved checkpoint checkpoints/checkpoint1.pt (epoch 1 @ 388 updates, score 6.286) (writing took 0.7026647419988876 seconds)
2021-12-06 00:50:59 | INFO | fairseq_cli.train | end of epoch 1 (average epoch stats below)                                                                               
2021-12-06 00:50:59 | INFO | train | epoch 001 | loss 7.141 | ppl 141.15 | wps 3143.8 | ups 0.31 | wpb 10178.1 | bsz 413 | num_updates 388 | lr 0.005 | gnorm 0.41 | train_wall 1232 | wall 1256
2021-12-06 00:50:59 | INFO | fairseq_cli.train | done training in 1255.9 seconds
$ 

無事終わりました。ちなみに学習に利用したデータはIWSLT'14 German to Englishです。独→英の機械翻訳タスクとなります。
こちらもfairseq/examples内にDL用のスクリプトが配置されているので、上記のリンクから飛んで実行してみると再現できます。

今回、独自定義以外で使用したオプションは以下となります。

オプション	意味
--arch	モデルアーキテクチャを指定。fairseq側でも定義済みのモノを利用できるし、独自定義も指定できる
--max-tokens	バッチ内の最大トークン数
--max-epoch	最大エポック数。強制的に指定したエポック数で学習を打ち切る
--optimizer	最適化関数を指定
--lr	学習率

推論をやってみる

学習が終了したら次はfairseq-generateコマンドを使用してテストセットに対して推論をやってみたいと思います。評価方法はBLEUスコアです。

$ fairseq-generate ../fairseq/data-bin/iwslt14.tokenized.de-en \
  --path checkpoints/checkpoint_best.pt \
  --beam 5 \
  --remove-bpe
...
P-7     -2.3089 -0.3036 -2.6581 -3.5390 -3.2376 -3.1657 -3.2949 -1.3530 -2.9390 -1.2448 -2.6206 -0.5986 -2.6413 -4.4828 -3.1905 -3.3060 -1.3208 -2.9843 -1.3485 -2.6342 -0.6018 -2.6975 -4.4914 -3.2837 -3.2641 -1.3087 -2.9723 -2.6705 -0.2006
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████                                                                                                                                                                          2021-12-06 12:48:27 | INFO | fairseq_cli.generate | NOTE: hypothesis and token scores are output in base 2
2021-12-06 12:48:27 | INFO | fairseq_cli.generate | Translated 6,750 sentences (118,941 tokens) in 1587.6s (4.25 sentences/s, 74.92 tokens/s)
Generate test with beam=5: BLEU4 = 2.69, 26.6/5.6/1.6/0.4 (BP=0.844, ratio=0.855, syslen=112131, reflen=131161)
$

ちゃんと動いていますね！
LSTMやGRUなどの再帰構造を持っているRNNは、その構造故に学習や推論が遅いと一般的に知られています。そこでfairseqは、以前の状態をキャッシュすることで高速化を実現しています。

以下のコードのincremental_state系がそちらの動作を実現しています。

from fairseq.models import FairseqIncrementalDecoder


class SimpleGRUDecoder(FairseqIncrementalDecoder):

    def __init__(
            self, dictionary, embed_dim,
            encoder_hidden=128, decoder_hidden=128, n_layer=1, dropout=0.5, w=None):
        super().__init__(dictionary)

        # decoder embedding
        self.decoder_embed = nn.Embedding(
            len(dictionary), embed_dim, padding_idx=dictionary.pad()
        )

        # 事前学習済み言語モデルがある場合はWを固定
        if w is not None:
            self.decoder_embed.weight.data.copy_(w)
            self.decoder_embed.requireds_grad_ = False

        # decoder
        self.decoder = nn.GRU(
            input_size=encoder_hidden + decoder_hidden,
            hidden_size=decoder_hidden,  # 隠れ層
            dropout=dropout,
            num_layers=n_layer,
            bidirectional=False
        )

        # dropout
        self.dropout = nn.Dropout(dropout)

        # 出力層
        self.output = nn.Linear(decoder_hidden, len(dictionary))

    def forward(
            self, prev_output_tokens, encoder_out, incremental_state=None):

        if incremental_state is not None:
            prev_output_tokens = prev_output_tokens[:, -1:]

        bsz, tgt_len = prev_output_tokens.size()
        final_encoder_hidden = encoder_out['final_hidden']

        # decoder embedding
        embed_trg = self.dropout(self.decoder_embed(prev_output_tokens))
        decoder_input = torch.cat(
            [embed_trg, final_encoder_hidden.unsqueeze(1).expand(bsz, tgt_len, -1)], dim=2,
        )

        initial_state = utils.get_incremental_state(
            self, incremental_state, 'prev_state')
        if initial_state is None:
            initial_state = final_encoder_hidden.unsqueeze(0)

        # decoder
        dec_output, latest_state = self.decoder(
            decoder_input.transpose(0, 1),
            initial_state)  # decoder output

        utils.set_incremental_state(
            self, incremental_state, 'prev_state', latest_state,
        )

        decoder_input = dec_output.transpose(0, 1)
        decoder_input = self.output(decoder_input)

        return decoder_input, None

    def reorder_incremental_state(self, incremental_state, new_order):

        prev_state = utils.get_incremental_state(
            self, incremental_state, 'prev_state',
        )

        # Reorder batches according to *new_order*.
        reordered_state = prev_state.index_select(1, new_order)  # hidden

        # Update the cached state.
        utils.set_incremental_state(
            self, incremental_state, 'prev_state', reordered_state,
        )

キャッシュ利用する形式に改良したモデルで推論を実施してみたいと思います。

$ fairseq-generate ../fairseq/data-bin/iwslt14.tokenized.de-en \
  --path checkpoints/checkpoint_best.pt \
  --beam 5 \
  --remove-bpe
…
P-7     -2.3089 -0.3036 -2.6581 -3.5390 -3.2376 -3.1657 -3.2949 -1.3530 -2.9390 -1.2448 -2.6206 -0.5986 -2.6413 -4.4828 -3.1905 -3.3060 -1.3208 -2.9843 -1.3485 -2.6342 -0.6018 -2.6975 -4.4914 -3.2837 -3.2641 -1.3087 -2.9723 -2.6705 -0.2006
2021-12-06 20:15:46 | INFO | fairseq_cli.generate | NOTE: hypothesis and token scores are output in base 2                                                                                                                                 
2021-12-06 20:15:46 | INFO | fairseq_cli.generate | Translated 6,750 sentences (118,941 tokens) in 72.4s (93.27 sentences/s, 1643.53 tokens/s)
Generate test with beam=5: BLEU4 = 2.69, 26.6/5.6/1.6/0.4 (BP=0.844, ratio=0.855, syslen=112131, reflen=131161)
$ 

キャッシュ	速度
なし	Translated 6,750 sentences (118,941 tokens) in 1587.6s (4.25 sentences/s, 74.92 tokens/s)
あり	Translated 6,750 sentences (118,941 tokens) in 72.4s (93.27 sentences/s, 1643.53 tokens/s)

え、早い…（ドン引き）

おわりに

正直、fairseqを見たときは「また変なライブラリ出てる。覚えるの辛い。。。」となり拒絶反応のためか記憶から消し去ろうとしましたが、実際使ってみるとこれまでの資産をほぼそのまま使えることやこれまで長々と実装していた部分をfairseq側に任せれることがわかりました。また推論も含め爆速化することが可能だとわかりました。

今後、Sequentialなモデルの学習・評価をする際のデファクトスタンダードツールとなりそうです。すでになっているかもしれません。

今回はfairseqで対応しているデータセットであったためデータ周りで苦労しませんでしたが、今後はオリジナルのデータセットを使用して学習を回したいと思った所存です。