nerman: AllenNLP と Optuna で作る固有表現抽出システム

事業開発部の @himkt です.好きなニューラルネットは BiLSTM-CRF です. 普段はクックパッドアプリのつくれぽ検索機能の開発チームで自然言語処理をしています.

本稿では,レシピテキストからの料理用語抽出システム nerman について紹介します. nerman の由来は ner (固有表現抽出 = Named Entity Recognition) + man (する太郎) です. クックパッドに投稿されたレシピから料理に関する用語を自動抽出するシステムであり,AllenNLP と Optuna を組み合わせて作られています. (コードについてすべてを説明するのは難しいため,実際のコードを簡略化している箇所があります)

料理用語の自動抽出

料理レシピには様々な料理用語が出現します. 食材や調理器具はもちろん,調理動作や食材の分量なども料理用語とみなせます. 「切る」という調理動作を考えても,「ざく切りにする」「輪切りにする」「みじん切りにする」など,用途に合わせて色々な切り方が存在します. レシピの中からこのような料理用語を抽出できれば,レシピからの情報抽出や質問応答などのタスクに応用できます.

料理用語の自動抽出には,今回は機械学習を利用します. 自然言語処理のタスクの中に,固有表現抽出というタスクが存在します. 固有表現抽出とは,自然言語の文(新聞記事などの文書が対象となることが多いです)から人名や地名,組織名などの固有表現を抽出するタスクです. このタスクは系列ラベリングと呼ばれる問題に定式化できます. 系列ラベリングを用いた固有表現抽出では,入力文を単語に分割したのち各単語に固有表現タグを付与します. タグが付与された単語列を抽出することで固有表現が得られます.

固有表現抽出の例
一般的な固有表現抽出(上)と料理用語への応用例(下)

今回は人名,地名などの代わりに食材名,調理器具名,調理動作の名前などを固有表現とみなしてモデルを学習します. 詳細な固有表現タグの定義は次の章で説明します.

データセット

機械学習モデルの学習には教師データが必要です. クックパッドでは言語データ作成の専門家の方に協力していただき,アノテーションガイドラインの整備およびコーパスの構築に取り組みました. レシピからの固有表現抽出については京都大学の森研究室でも研究されています(論文はこちら. PDF ファイルが開かれます). この研究で定義されている固有表現タグを参考にしつつ,クックパッドでのユースケースに合わせて次のような固有表現タグを抽出対象として定義しました.

固有表現タグの一覧
固有表現タグの一覧

この定義に基づき,クックパッドに投稿されたレシピの中から 500 品のレシピに対して固有表現を付与しました. データは Cookpad Parsed Corpus と名付けられ,社内の GitHub リポジトリで管理されています. また,機械学習モデルで利用するための前処理(フォーマットの変更など)をしたデータが S3 にアップロードされています.

Cookpad Parsed Corpus に関するアウトプットとして論文化にも取り組んでいます. 執筆した論文は自然言語処理の国際会議である COLING で開催される言語資源に関する研究のワークショップ LAW(Linguistic Annotation Workshop)に採択されました. 🎊

タイトルは以下の通りです.

Cookpad Parsed Corpus: Linguistic Annotations of Japanese Recipes
Jun Harashima and Makoto Hiramatsu

Cookpad Parsed Corpus に収録されているレシピは固有表現の他にも形態素と係り受けの情報が付与されており, 現在大学等の研究機関に所属されている方に利用いただけるように公開の準備を進めています.

準備: AllenNLP を用いた固有表現抽出モデル

nerman ではモデルは AllenNLP を用いて実装しています.

github.com

AllenNLP は Allen Institute for Artificial Intelligence (AllenAI) が開発している自然言語処理フレームワークであり, 最新の機械学習手法に基づく自然言語処理のためのニューラルネットワークを簡単に作成できる便利なライブラリです. AllenNLP は pip でインストールできます.

pip install allennlp

AllenNLP ではモデルの定義や学習の設定を Jsonnet 形式のファイルに記述します. 以下に今回の固有表現抽出モデルの学習で利用する設定ファイル(config.jsonnet)を示します. (モデルは BiLSTM-CRF を採用しています.)

  • config.jsonnet
local dataset_base_url = 's3://xxx/nerman/data';

{
  dataset_reader: {
    type: 'cookpad2020',
    token_indexers: {
      word: {
        type: 'single_id',
      },
    },
    coding_scheme: 'BIOUL',
  },
  train_data_path: dataset_base_url + '/cpc.bio.train',
  validation_data_path: dataset_base_url + '/cpc.bio.dev',
  model: {
    type: 'crf_tagger',
    text_field_embedder: {
      type: 'basic',
      token_embedders: {
        word: {
          type: 'embedding',
          embedding_dim: 32,
        },
      },
    },
    encoder: {
      type: 'lstm',
      input_size: 32,
      hidden_size: 32,
      dropout: 0.5,
      bidirectional: true,
    },
    label_encoding: 'BIOUL',
    calculate_span_f1: true,
    dropout: 0.5,
    initializer: {},
  },
  data_loader: {
    batch_size: 10,
  },
  trainer: {
    num_epochs: 20,
    cuda_device: -1,
    optimizer: {
      type: 'adam',
      lr: 5e-4,
    },
  },
}

モデル,データ,そして学習に関する設定がそれぞれ指定されています. AllenNLP はデータセットのパスとしてローカルのファイルパスだけでなく URL を指定できます. 現状では httphttps,そして s3 のスキーマに対応しているようです. (読んだコードはこのあたり) nerman では train_data_path および validation_data_path に S3 上の加工済み Cookpad Parsed Corpus の学習データ,バリデーションデータの URL を指定しています.

AllenNLP は自然言語処理の有名なタスクのデータセットを読み込むためのコンポーネントを提供してくれます. しかしながら,今回のように自分で構築したデータセットを利用したい場合には自分でデータセットをパースするクラス(データセットリーダー)を作成する必要があります. cookpad2020 は Cookpad Parsed Corpus を読み込むためのデータセットリーダーです. データセットリーダーの作成方法については公式チュートリアルで 説明されているので詳しく知りたい方はそちらを参照いただければと思います.

設定ファイルが作成できたら, allennlp train config.jsonnet --serialization-dir result のようにコマンドを実行することで学習がはじまります. 学習のために必要な情報すべてが設定ファイルにまとまっていて,実験を管理しやすいことが AllenNLP の特徴の1つです. serialization-dir については後述します.

今回の記事では紹介しませんが, allennlp コマンドには allennlp predict allennlp evaluate などの非常に便利なサブコマンドが用意されています. 詳しく知りたい方は公式ドキュメントを参照ください.

nerman の全体像

以下に nerman の全体像を示します.

nerman の全体像
nerman の全体像

システムは大きく分けて 3 つのバッチから構成されています.それぞれの役割は以下の通りです.

  • (1) ハイパーパラメータ最適化
  • (2) モデルの学習
  • (3) 実データ(レシピ)からの固有表現抽出(予測)

本稿では,順序を入れ替えて モデルの学習 => 実データでの予測 => ハイパーパラメータ最適化 の順に解説していきます.

モデルの学習

モデルの学習バッチは以下のようなシェルスクリプトを実行します.

  • train
#!/bin/bash

allennlp train \
  config/ner.jsonnet \
  --serialization-dir result \
  --include-package nerman

# モデルとメトリクスのアップロード
aws s3 cp result/model.tar.gz s3://xxx/nerman/model/$TIMESTAMP/model.tar.gz
aws s3 cp result/metrics.json s3://xxx/nerman/model/$TIMESTAMP/metrics.json

準備の章で解説したように, allennlp train コマンドでモデルを学習します. --serialization-dir で指定しているディレクトリにはモデルのアーカイブ(tar.gz 形式), アーカイブファイルの中にはモデルの重みの他に標準出力・標準エラー出力,そして学習したモデルのメトリクスなどのデータが保存されます.

学習が終わったら, allennlp train コマンドによって生成されたモデルのアーカイブとメトリクスを S3 にアップロードします. (アーカイブファイルにはモデルの重みなどが保存されており,このファイルがあれば即座にモデルを復元できます.) また,メトリクスファイルも同時にアップロードしておくことで,モデルの性能をトラッキングできます.

S3
S3 の様子.実行日ごとにモデルのアーカイブとメトリクスがアップロードされる.

  • metrics.json

生成されるメトリクスファイル.性能指標だけでなく学習時間や計算時間などもわかります)

{
  "best_epoch": 19,
  "peak_worker_0_memory_MB": 431.796,
  "training_duration": "0:29:38.785065",
  "training_start_epoch": 0,
  "training_epochs": 19,
  "epoch": 19,
  "training_accuracy": 0.8916963871929718,
  "training_accuracy3": 0.8938523846944327,
  "training_precision-overall": 0.8442808607021518,
  "training_recall-overall": 0.8352005377548734,
  "training_f1-measure-overall": 0.8397161522865011,
  "training_loss": 38.08172739275527,
  "training_reg_loss": 0.0,
  "training_worker_0_memory_MB": 431.796,
  "validation_accuracy": 0.8663015463917526,
  "validation_accuracy3": 0.8688788659793815,
  "validation_precision-overall": 0.8324965769055226,
  "validation_recall-overall": 0.7985989492119089,
  "validation_f1-measure-overall": 0.815195530726207,
  "validation_loss": 49.37634348869324,
  "validation_reg_loss": 0.0,
  "best_validation_accuracy": 0.8663015463917526,
  "best_validation_accuracy3": 0.8688788659793815,
  "best_validation_precision-overall": 0.8324965769055226,
  "best_validation_recall-overall": 0.7985989492119089,
  "best_validation_f1-measure-overall": 0.815195530726207,
  "best_validation_loss": 49.37634348869324,
  "best_validation_reg_loss": 0.0,
  "test_accuracy": 0.875257568552861,
  "test_accuracy3": 0.8789031542241242,
  "test_precision-overall": 0.8318906605922551,
  "test_recall-overall": 0.8214125056230319,
  "test_f1-measure-overall": 0.8266183793571253,
  "test_loss": 48.40180677297164
}

モデルの学習は EC2 インスタンス上で実行されます. 今回のケースではデータセットは比較的小さく(全データ = 500 レシピ), BiLSTM-CRF のネットワークもそこまで大きくありません. このため,通常のバッチジョブとほぼ同じ程度の規模のインスタンスでの学習が可能です. 実行環境が GPU や大容量メモリなどのリソースを必要としないため,通常のバッチ開発のフローに乗ることができました. これにより,社内に蓄積されていたバッチ運用の知見を活かしてインフラ環境の整備にかかるコストを抑えつつ学習バッチを構築できています.

また, nerman のバッチはすべてスポットインスタンスを前提として構築されています. スポットインスタンスは通常のインスタンスよりもコストが低く, 代わりに実行中に強制終了する(spot interruption と呼ばれる)可能性があるインスタンスです. モデルの学習は強制終了されてしまってもリトライをかければよく,学習にかかる時間が長すぎなければスポットインスタンスを利用することでコストを抑えられます. (ただし,学習にかかる時間が長ければ長いだけ spot interruption に遭遇する可能性が高くなります. リトライを含めた全体での実行時間が通常のインスタンスでの実行時間と比較して長くなりすぎた場合, かえってコストがかかってしまう可能性があり,注意が必要です.)

実データでの予測

以下のようなシェルスクリプトを実行して予測を実行します.

  • predict
#!/bin/bash

export MODEL_VERSION=${MODEL_VERSION:-2020-07-08}
export TIMESTAMP=${TIMESTAMP:-`date '+%Y-%m-%d'`}

export FROM_IDX=${FROM_IDX:-10000}
export LAST_IDX=${LAST_IDX:-10100}

export KUROKO2_PARALLEL_FORK_INDEX=${KUROKO2_PARALLEL_FORK_INDEX:--1}
export KUROKO2_PARALLEL_FORK_SIZE=${KUROKO2_PARALLEL_FORK_SIZE:--1}

if [ $KUROKO2_PARALLEL_FORK_SIZE = -1 ] || [ $KUROKO2_PARALLEL_FORK_INDEX = -1 ]; then
  echo $FROM_IDX $LAST_IDX ' (without parallel execution)'

else
  if (($KUROKO2_PARALLEL_FORK_INDEX >= $KUROKO2_PARALLEL_FORK_SIZE)); then
    echo '$KUROKO2_PARALLEL_FORK_INDEX'=$KUROKO2_PARALLEL_FORK_INDEX 'must be smaller than $KUROKO2_PARALLEL_FORK_SIZE' $KUROKO2_PARALLEL_FORK_SIZE
    exit
  fi

  # ==============================================================================
  # begin: FROM_IDX ~ LAST_IDX のデータを KUROKO2_PARALLEL_FORK_SIZE の値で等分する処理
  # ==============================================================================

  NUM_RECORDS=$(($LAST_IDX - $FROM_IDX))
  echo 'NUM_RECORDS = ' $NUM_RECORDS

  if (($NUM_RECORDS % $KUROKO2_PARALLEL_FORK_SIZE != 0)); then
    echo '$KUROKO2_PARALLEL_FORK_SIZE = ' $KUROKO2_PARALLEL_FORK_SIZE 'must be multiple of $NUM_RECORDS=' $NUM_RECORDS
    exit
  fi

  DIV=$(($NUM_RECORDS / $KUROKO2_PARALLEL_FORK_SIZE))
  echo 'DIV=' $DIV

  if (($DIV <= 0)); then
    echo 'Invalid DIV=' $DIV
    exit
  fi

  LAST_IDX=$(($FROM_IDX + (($KUROKO2_PARALLEL_FORK_INDEX + 1) * $DIV) ))
  FROM_IDX=$(($FROM_IDX + ($KUROKO2_PARALLEL_FORK_INDEX * $DIV) ))
  echo '$FROM_IDX = ' $FROM_IDX ' $LAST_IDX = ' $LAST_IDX

  # ============================================================================
  # end: FROM_IDX ~ LAST_IDX のデータを KUROKO2_PARALLEL_FORK_SIZE の値で等分する処理
  # ============================================================================
fi

allennlp custom-predict \
    --from-idx $FROM_IDX \
    --last-idx $LAST_IDX \
    --include-package nerman \
    --model-path s3://xxx/nerman/model/$MODEL_VERSION/model.tar.gz

aws s3 cp \
    --recursive \
    --exclude "*" \
    --include "_*.csv" \
    prediction \
    s3://xxx/nerman/output/$TIMESTAMP/prediction/

予測バッチは学習バッチが作成したモデルを読み込み,固有表現が付与されていないレシピを解析します. また,予測バッチは並列実行に対応しています. クックパッドには 340 万品以上のレシピが投稿されており,これらのレシピを一度に解析するのは容易ではありません. このため,レシピを複数のグループに分割し,それぞれを並列に解析しています.

並列処理の様子
並列処理の様子

FROM_RECIPE_IDXLAST_RECIPE_IDX で解析対象とするレシピを指定し, KUROKO2_PARALLEL_FORK_SIZE という環境変数で並列数を設定します. 並列実行されたプロセスには KUROKO2_PARALLEL_FORK_INDEX という変数が渡されるようになっていて,この変数で自身が並列実行されたプロセスのうち何番目かを識別します. プロセスの並列化は社内で利用されているジョブ管理システム kuroko2 の並列実行機能 (parallel_fork) を利用して実現しています.

custom-predict コマンドは上で定義した変数を用いて対象となるレシピを分割し, AllenNLP のモデルを用いて固有表現を抽出するコマンドです. AllenNLP では自分でサブコマンドを登録でき,このようにすべての処理を allennlp コマンドから実行できるようになっています. サブコマンドは以下のように Python スクリプト(predict_custom.py)を作成して定義できます. (サブコマンドについての公式ドキュメントはこちら

  • custom_predict.py
import argparse

from allennlp.commands import Subcommand

from nerman.data.dataset_readers import StreamSentenceDatasetReader
from nerman.predictors import KonohaSentenceTaggerPredictor


def create_predictor(model_path) -> KonohaSentenceTaggerPredictor:
    archive = load_archive(model_path)
    predictor = KonohaSentenceTaggerPredictor.from_archive(archive)
    dataset_reader = StreamSentenceDatasetReader(predictor._dataset_reader._token_indexers)
    return KonohaSentenceTaggerPredictor(predictor._model, dataset_reader)


def _predict(
  from_idx: int,
  last_idx: int,
  model_path: str,
):

    # predictor の作成
    predictor = create_predictor(model_path)
    ...  # Redshift からデータを取ってきたりモデルに入力したりする処理(今回の記事では解説は割愛します)


def predict(args: argparse.Namespace):
    from_idx = args.from_idx
    last_idx = args.last_idx
    _predict(from_idx, last_idx)


@Subcommand.register("custom-predict")
class CustomPrediction(Subcommand):
    @overrides
    def add_subparser(self, parser: argparse._SubParsersAction) -> argparse.ArgumentParser:
        description = "Script to custom predict."
        subparser = parser.add_parser(self.name, description=description, help="Predict entities.")

        subparser.add_argument("--from-idx", type=int, required=True)
        subparser.add_argument("--last-idx", type=int, required=True)
        subparser.add_argument("--model-path", type=str, required=True)

        subparser.set_defaults(func=predict)  # サブコマンドが呼ばれたときに実際に実行するメソッドを指定する
        return subparser

model_path という変数にはモデルのアーカイブファイルのパスが指定されています. アーカイブファイルのパスは load_archive というメソッドに渡されます. load_archive は AllenNLP が提供しているメソッドであり,これを利用すると保存された学習済みモデルの復元が簡単にできます. また, load_archive はデータセットのパスと同様 S3 スキーマに対応しているため,学習バッチでアップロード先に指定したパスをそのまま利用できます. (load_archive の公式ドキュメントはこちら

文字列をモデルに入力するためには AllenNLP の Predictor という機構を利用しています. 公式ドキュメントはこちらです. 系列ラベリングモデルの予測結果を扱う際に便利な SentenceTaggerPredictor クラスを継承し,以下に示す KonohaSentenceTaggerPredictor クラスを定義しています. predict メソッドに解析したい文字列を入力すると,モデルの予測結果を出力してくれます.

from allennlp.common.util import JsonDict
from allennlp.data import Instance
from allennlp.data.dataset_readers.dataset_reader import DatasetReader
from allennlp.models import Model
from allennlp.predictors import SentenceTaggerPredictor
from allennlp.predictors.predictor import Predictor
from konoha.integrations.allennlp import KonohaTokenizer
from overrides import overrides


@Predictor.register("konoha_sentence_tagger")
class KonohaSentenceTaggerPredictor(SentenceTaggerPredictor):
    def __init__(self, model: Model, dataset_reader: DatasetReader) -> None:
        super().__init__(model, dataset_reader)
        self._tokenizer = KonohaTokenizer("mecab")

    def predict(self, sentence: str) -> JsonDict:
        return self.predict_json({"sentence": sentence})

    @overrides
    def _json_to_instance(self, json_dict: JsonDict) -> Instance:
        sentence = json_dict["sentence"]
        tokens = self._tokenizer.tokenize(sentence)
        return self._dataset_reader.text_to_instance(tokens)

nerman では,日本語のレシピデータを扱うために日本語処理ツールの konoha を利用しています. KonohaTokenizer は Konoha が提供している AllenNLP インテグレーション機能です. 日本語文字列を受け取り,分かち書きもしくは形態素解析を実施, AllenNLP のトークン列を出力します. 形態素解析器には MeCab を採用しており,辞書は mecab-ipadic を使用しています.

github.com

次に,作成したモジュールを __init__.py でインポートします. 今回は nerman/commands というディレクトリに custom_predict.py を設置しています. このため, nerman/__init__.py および nerman/commands/__init__.py をそれぞれ次のように編集します.

  • nerman/__init__.py
import nerman.commands
  • nerman/commands/__init__.py
from nerman.commands import custom_predict

コマンドの定義およびインポートができたら, allennlp コマンドで実際にサブコマンドを認識させるために .allennlp_plugins というファイルをリポジトリルートに作成します.

  • .allennlp_plugins
nerman

以上の操作でサブコマンドが allennlp コマンドで実行できるようになります. allennlp --help を実行して作成したコマンドが利用できるようになっているか確認できます.

得られた予測結果は CSV 形式のファイルとして保存され,予測が終了した後に S3 へアップロードされます.

次に, S3 にアップロードした予測結果をデータベースに投入します. データは最終的に Amazon Redshift (以降 Redshift) に配置されますが, Amazon Aurora (以降 Aurora)を経由するアーキテクチャを採用しています. これは Aurora の LOAD DATA FROM S3 ステートメントという機能を利用するためです. LOAD DATA FROM S3 ステートメントは次のような SQL クエリで利用できます.

  • load.sql
load
    data
from
    S3 's3://xxx/nerman/output/$TIMESTAMP/prediction.csv'
into table recipe_step_named_entities
    fields terminated by ','
    lines  terminated by '\n'
    (recipe_text_id, start, last, name, category)
    set created_at = current_timestamp, updated_at = current_timestamp;

このクエリを実行することで, S3 にアップロードした CSV ファイルを直接 Amazon Aurora にインポートできます. LOAD DATA FROM S3 については AWS の公式ドキュメント が参考になります. バッチサイズやコミットのタイミングの調整の手間が必要なくなるため,大規模データをデータベースに投入する際に非常に便利です.

Aurora のデータベースに投入した予測結果は pipelined-migrator という社内システムを利用して定期的に Redshift へ取り込まれます. pipelined-migrator を利用することで,管理画面上で数ステップ設定を行うだけで Aurora から Redshift へデータを取り込めます. これにより, S3 からのロードと pipelined-migrator を組み合わせた手間の少ないデータの投入フローが実現できました.

解析結果をスタッフに利用してもらう方法として,データベースを利用せずに予測 API を用意する方法も考えられます. 今回のタスクの目標は「すでに投稿されたレシピからの料理用語の自動抽出」であり,これはバッチ処理であらかじめ計算可能です. このため, API サーバを用意せずにバッチ処理で予測を行う方針を採用しました.

また,エンジニア以外のスタッフにも予測結果を使ってみてもらいたいと考えていました. クックパッドはエンジニア以外のスタッフも SQL を書ける方が多いため, 予測結果をクエリ可能な形でデータベースに保存しておく方針はコストパフォーマンスがよい選択肢でした. 予測結果を利用するクエリ例を以下に示します.

  • list_tools.sql
select
    , s.recipe_id
    , e.name
    , e.category
from
    recipe_step_named_entities as e
    inner join recipe_steps as s on e.step_id = s.id
where
    e.category in ('Tg')
    and s.recipe_id = xxxx

このクエリを Redshift 上で実行することで,レシピ中に出現する調理器具のリストを取得できるようになりました.

SQL の実行結果
SQL の実行結果

Optuna を用いたハイパーパラメータの分散最適化

最後にハイパーパラメータの最適化について解説します.

github.com

nerman では Optuna を用いたハイパーパラメータの最適化を実施しています. Optuna は Preferred Networks (PFN) が開発しているハイパーパラメータ最適化のライブラリです. インストールは pip install optuna をターミナルで実行すれば完了します.

Optuna では,各インスタンスから接続可能なバックエンドエンジン(RDB or Redis)を用意し,それをストレージで使用することで, 複数インスタンスを利用した分散環境下でのハイパーパラメータ最適化を実現できます. (ストレージは Optuna が最適化結果を保存するために使用するもので,RDB や Redis などを抽象化したものです) インスタンスをまたいだ分散最適化を実施する場合,ストレージのバックエンドエンジンは MySQL もしくは PostgreSQL が推奨されています (Redis も experimental feature として利用可能になっています). 詳しくは公式ドキュメントをご参照ください. 今回はストレージとして MySQL (Aurora) を採用しています.

Optuna には AllenNLP のためのインテグレーションモジュールが存在します. しかしながら,このインテグレーションモジュールを使うと自身で最適化を実行するための Python スクリプトを記述する必要があります. そこで, AllenNLP とよりスムーズに連携するために allennlp-optuna というツールを開発しました. allennlp-optuna をインストールすると,ユーザは allennlp tune というコマンドで Optuna を利用したハイパーパラメータ最適化を実行できるようになります. このコマンドは allennlp train コマンドと互換性が高く, AllenNLP に慣れたユーザはスムーズにハイパーパラメータの最適化を試せます.

github.com

allennlp tune コマンドを実行するには,まず pip install allennlp-optuna.gitallennlp-optuna をインストールします. 次に, .allennlp_plugins を以下のように編集します.

  • .allennlp_plugins
allennlp-optuna
nerman

allennlp --help とコマンドを実行して,以下のように retrain コマンドと tune コマンドが確認できればインストール成功です.

$ allennlp --help
2020-11-05 01:54:24,567 - INFO - allennlp.common.plugins - Plugin allennlp_optuna available
usage: allennlp [-h] [--version]  ...

Run AllenNLP

optional arguments:
  -h, --help     show this help message and exit
  --version      show program's version number and exit

Commands:

    best-params  Export best hyperparameters.
    evaluate     Evaluate the specified model + dataset.
    find-lr      Find a learning rate range.
    predict      Use a trained model to make predictions.
    print-results
                 Print results from allennlp serialization directories to the console.
    retrain      Train a model with hyperparameter found by Optuna.
    test-install
                 Test AllenNLP installation.
    train        Train a model.
    tune         Optimize hyperparameter of a model.

allennlp-optuna が無事にインストールできました. 次に allennlp-optuna を利用するために必要な準備について解説します.

設定ファイルの修正

はじめに,準備の章で作成した config.jsonnet を以下のように書き換えます.

  • config.jsonnet (allennlp-optuna 用)
// ハイパーパラメータを変数化する
local lr = std.parseJson(std.extVar('lr'));
local lstm_hidden_size = std.parseInt(std.extVar('lstm_hidden_size'));
local dropout = std.parseJson(std.extVar('dropout'));
local word_embedding_dim = std.parseInt(std.extVar('word_embedding_dim'));

local cuda_device = -1;

{
  dataset_reader: {
    type: 'cookpad2020',
    token_indexers: {
      word: {
        type: 'single_id',
      },
    },
    coding_scheme: 'BIOUL',
  },
  train_data_path: 'data/cpc.bio.train',
  validation_data_path: 'data/cpc.bio.dev',
  model: {
    type: 'crf_tagger',
    text_field_embedder: {
      type: 'basic',
      token_embedders: {
        word: {
          type: 'embedding',
          embedding_dim: word_embedding_dim,
        },
      },
    },
    encoder: {
      type: 'lstm',
      input_size: word_embedding_dim,
      hidden_size: lstm_hidden_size,
      dropout: dropout,
      bidirectional: true,
    },
    label_encoding: 'BIOUL',
    calculate_span_f1: true,
    dropout: dropout,  // ここで宣言した変数を指定する
    initializer: {},
  },
  data_loader: {
    batch_size: 10,
  },
  trainer: {
    num_epochs: 20,
    cuda_device: cuda_device,
    optimizer: {
      type: 'adam',
      lr: lr,  // ここで宣言した変数を指定する
    },
  },
}

最適化したいハイパーパラメータを local lr = std.parseJson(std.extVar('lr')) のように変数化しています. std.extVar の返り値は文字列です.機械学習モデルのハイパーパラメータは整数や浮動小数であることが多いため,キャストが必要になります. 浮動小数へのキャストは std.parseJson というメソッドを利用します.整数へのキャストは std.parseInt を利用してください.

探索空間の定義

次に,ハイパーパラメータの探索空間を定義します. allennlp-optuna では,探索空間は次のような JSON ファイル(hparams.json)で定義します.

  • hparams.json
[
  {
    "type": "float",
    "attributes": {
      "name": "dropout",
      "low": 0.0,
      "high": 0.8
    }
  },
  {
    "type": "int",
    "attributes": {
      "name": "lstm_hidden_size",
      "low": 32,
      "high": 256
    },
  },
  {
    "type": "float",
    "attributes": {
      "name": "lr",
      "low": 5e-3,
      "high": 5e-1,
      "log": true
    }
  }
]

今回の例では学習率とドロップアウトの比率が最適化の対象です. それぞれについて,値の上限・下限を設定します. 学習率は対数スケールの分布からサンプリングするため, "log": true としていることに注意してください.

最適化バッチは次のようなシェルスクリプトを実行します.

  • optimize
#!/bin/bash

export N_TRIALS=${N_TRIALS:-20}  # Optuna の試行回数を制御する
export TIMEOUT=${TIMEOUT:-7200}  # # 一定時間が経過したら最適化を終了する(単位は秒): 60*60*2 => 2h
export TIMESTAMP=${TIMESTAMP:-`date '+%Y-%m-%d'`}

export OPTUNA_STORAGE=${OPTUNA_STORAGE:-mysql://$DB_USERNAME:$DB_PASSWORD@$DB_HOST_NAME/$DB_NAME}
export OPTUNA_STUDY_NAME=${OPTUNA_STUDY_NAME:-nerman-$TIMESTAMP}

# ハイパーパラメータの最適化
allennlp tune \
  config/ner.jsonnet \
  config/hparam.json \
  --serialization-dir result/hpo \
  --include-package nerman \
  --metrics best_validation_f1-measure-overall \
  --study-name $OPTUNA_STUDY_NAME \
  --storage $OPTUNA_STORAGE \
  --direction maximize \
  --n-trials $N_TRIALS \
  --skip-if-exists \
  --timeout $TIMEOUT

このコマンドを複数のインスタンスで実行することで,ハイパーパラメータの分散最適化が実行できます. オプション --skip-if-exists を指定することで,複数のインスタンスの間で最適化の途中経過を共有しています. Optuna は通常実行のたびに新しく実験環境(study と呼ばれます)を作成し,ハイパーパラメータの探索を行います. このとき,すでにストレージに同名の study が存在する場合はエラーになります. しかし, --skip-if-exists を有効にすると,ストレージに同名の study がある場合は当該の study を読み込み,途中から探索を再開します. この仕組みによって,複数のインスタンスで --skip-if-exists を有効にして探索を開始することでだけで study を共有した最適化が行われます. 上記のスクリプトによって,最適化バッチは与えられた時間(--timeout で設定されている値 = 2 時間)に最大 20 回探索を実行します.

このように, Optuna のリモートストレージ機能によって,複数のインスタンスで同じコマンドを実行するだけで分散最適化が実現できました! Optuna の分散ハイパーパラメータ最適化の詳しい仕組み,あるいはより高度な使い方については Optuna 開発者の 解説資料 が参考になるので, 興味のある方は合わせてご参照ください.

モデルの再学習

最後に,最適化されたハイパーパラメータを用いてモデルを再学習します. 再学習バッチは以下のようなシェルスクリプトで実行します.

  • retrain
#!/bin/bash

export TIMESTAMP=${TIMESTAMP:-`date '+%Y-%m-%d'`}

export OPTUNA_STORAGE=${OPTUNA_STORAGE:-mysql://$DB_USERNAME:$DB_PASSWORD@$DB_HOST_NAME/$DB_NAME}

# 最適化されたハイパーパラメータを用いたモデルの再学習
allennlp retrain \
  config/ner.jsonnet \
  --include-package nerman \
  --include-package allennlp_models \
  --serialization-dir result \
  --study-name $OPTUNA_STUDY_NAME \
  --storage $OPTUNA_STORAGE

# モデルとメトリクスのアップロード
aws s3 cp result/model.tar.gz s3://xxx/nerman/model/$TIMESTAMP/model.tar.gz
aws s3 cp result/metrics.json s3://xxx/nerman/model/$TIMESTAMP/metrics.json

このシェルスクリプトでは allennlp-optuna が提供する retrain コマンドを利用しています. allennlp retrain コマンドはストレージから最適化結果を取得し,得られたハイパーパラメータを AllenNLP に渡してモデルの学習を行ってくれます. tune コマンド同様, retrain コマンドも train コマンドとほぼ同じインターフェースを提供していることがわかります.

再学習したモデルのメトリクスを以下に示します.

  • metrics.json
{
  "best_epoch": 2,
  "peak_worker_0_memory_MB": 475.304,
  "training_duration": "0:45:46.205781",
  "training_start_epoch": 0,
  "training_epochs": 19,
  "epoch": 19,
  "training_accuracy": 0.9903080859981059,
  "training_accuracy3": 0.9904289830542626,
  "training_precision-overall": 0.9844266427651112,
  "training_recall-overall": 0.9843714989917096,
  "training_f1-measure-overall": 0.9843990701061036,
  "training_loss": 3.0297666011196327,
  "training_reg_loss": 0.0,
  "training_worker_0_memory_MB": 475.304,
  "validation_accuracy": 0.9096327319587629,
  "validation_accuracy3": 0.911243556701031,
  "validation_precision-overall": 0.884530630233583,
  "validation_recall-overall": 0.8787215411558669,
  "validation_f1-measure-overall": 0.8816165165824231,
  "validation_loss": 61.33201414346695,
  "validation_reg_loss": 0.0,
  "best_validation_accuracy": 0.9028672680412371,
  "best_validation_accuracy3": 0.9048002577319587,
  "best_validation_precision-overall": 0.8804444444444445,
  "best_validation_recall-overall": 0.867338003502627,
  "best_validation_f1-measure-overall": 0.873842082046708,
  "best_validation_loss": 38.57948366800944,
  "best_validation_reg_loss": 0.0,
  "test_accuracy": 0.8887303851640513,
  "test_accuracy3": 0.8904739261372642,
  "test_precision-overall": 0.8570790531487271,
  "test_recall-overall": 0.8632478632478633,
  "test_f1-measure-overall": 0.8601523980277404,
  "test_loss": 44.22851959539919
}

モデルの学習 の章で学習されたモデルと比較して,テストデータでの F値(test_f1-measure-overall)が 82.7 から 86.0 となり, 3.3 ポイント性能が向上しました. ハイパーパラメータの探索空間をアバウトに定めて Optuna に最適化をしてもらえば十分な性能を発揮するハイパーパラメータが得られます.便利です.

Optuna はハイパーパラメータを最適化するだけでなく, 最適化途中のメトリクスの推移やハイパーパラメータの重要度などを可視化する機能, 最適化結果を pandas DataFrame で出力する機能をはじめとする強力な実験管理機能を提供しています. より詳しく AllenNLP と Optuna の使い方を学びたい方は AllenNLP の公式ガイド なども合わせて読んでみてください.

まとめ

本稿では AllenNLP と Optuna を用いて構築した固有表現抽出システム nerman について紹介しました. nerman は AllenNLP を用いたモデル学習・実データ適用, Amazon Aurora を活用したデータ投入の手間の削減, および Optuna を活用したスケーラブルなハイパーパラメータ探索を実現しています. AllenNLP と Optuna を用いた機械学習システムの一例として,読んでくださった皆さんの参考になればうれしいです.

クックパッドでは自然言語処理の技術で毎日の料理を楽しくする仲間を募集しています. 実現したい価値のため,データセットの構築から本気で取り組みたいと考えている方にはとても楽しめる環境だと思います. 興味をもってくださった方はぜひご応募ください! クックパッドの {R&D, サービス開発現場} での自然言語処理についてカジュアルに話を聞きたいと思ってくださった方は @himkt までお気軽にご連絡ください.