機械学習における技術的負債をDigdagで返済する

お久しぶりです。技術開発部の相原です。 昨年度は技術基盤部としてmrubyを導入したりしていましたが今は少しレイヤーが開発寄りになりました。

とはいえ依然として技術基盤も見ていて、最近はご多分に漏れず機械学習を用いた技術基盤の改善に興味があります。 そんな中でここ数ヶ月メインの業務の合間の時間を使って試験的に機械学習を導入していたので、今回は技術的負債の高利子クレジットカードと呼ばれる機械学習を導入する中でどのような工夫をしたかということについて書きたいと思います。

機械学習については門外漢なので、ここではモデルの訓練などのプラクティスに関しては触れません。 (一部暗黙的に深層学習を前提としている箇所がありますのでご了承ください)

技術的負債の高利子クレジットカード

機械学習は運用における課題が多いことから技術的負債の高利子クレジットカードと呼ばれることがあります。

初出は恐らくGoogleが2014年のNIPSで発表した論文Machine Learning: The High Interest Credit Card of Technical Debtで、論文中では以下の要因によって技術的負債が生み出されるとされています。

1. 複雑なモデルが外部データに依存することによるシステムの抽象境界の侵食(Complex Models Erode Boundaries)
2. コードへの依存性よりも強いデータへの依存性(Data Dependencies Cost More than Code Dependencies)
3. 複雑なパイプラインの乱立に代表されるシステムのスパゲッティ化(System-level Spaghetti)
4. 不安定な外部世界への依存(Dealing with Changes in the External World)

今回はこれらの課題を解決して安定運用を実現すべく、この中のシステム面に着目しました。

Data Dependencies Cost More than Code Dependencies

データへの依存性が強いことに関する問題として、入力として与えられるデータが不安定であることによるモデルへの影響が挙げられます。 入力値となるデータの一部が欠損している可能性やそもそもデータの品質に問題があるケースが考えられ、またデータが時の変化によって活用されなくなったり精度に寄与しないことが後から分かることもあります。

そのためシステムとしてこれらの措置を講じることで不安定なデータによるモデルへの影響を検知できるようにする必要があります。

1. 入力データの静的解析
2. 入力データのバージョン管理
3. 不要な特徴量を取り除くための継続的な精度測定

System-level Spaghetti

機械学習を用いたシステムが陥りやすい問題としてグルーコードデザインパターンというものが存在します。 これは一つのシステムで多くの問題を解決しようとした際に発生しやすく、データの前処理に係る複雑なパイプラインの乱立によってシステムの大部分をグルーコードが占めることとなります。 こうしたグルーコードデザインパターンはシステムのメンテナンスを困難とし、時として再実装が必要となるケースもあります。

また機械学習を用いたシステムは往々にして膨大な設定を必要としますが、これらの設定はモデルに大きな影響を及ぼすものの、コードに関するテストは行われやすい一方で設定のテストは軽視されがちな傾向もあります。

他にも構築したシステムを自立させるためにはタスクのスケジューリングも必要となり、更に転移学習を適用するとなるとそれらのスケジューリングにも依存関係を持たせる必要が出てくることを考慮する必要もあるでしょう。

そのためシステムとしては以下のような要件が求められることになります。

1. パイプラインの責務の分離
2. 設定値のバージョン管理とテスタブルな設計
3. 各タスクのスケジューリングと依存管理

Dealing with Changes in the External World

機械学習を用いたシステムは外部世界の変化に直接影響される事が多くなります。 外部世界が安定することは稀であり、複数の特徴量から学習していた場合にそれらの特徴量に相関がなくなりモデルに影響を及ぼすといったケースが考えられます。

外部世界の変化によるモデルへの影響を検知するため、システムには以下が求められることになります。

1. 精度のモニタリングとアラート
2. サニティチェックとしての異常値の検出

Digdagによるアプローチ

ここからはそれぞれの課題についてDigdagを用いてどのようなアプローチを取ったかということを紹介していきます。

Digdagとは

まずはじめにDigdagとはオープンソースのワークフローエンジンです。 .dig というyamlによく似た設定ファイルによって各ジョブの順序や並列度、リトライ処理や処理に失敗した際のアラート処理を定義することができます。 サーバモードで起動することによって常駐してワークフローの定期実行のスケジューリングを行うことも可能で、その際にワークフロー間の依存関係を持たせることも可能です。

今回は機械学習における複雑なワークフローに対してこのDigdagを適用しました。

実際にワークフローを構築してみる

これまでのシステムに求められる要件をまとめると以下のようになります。

1. 入力データの静的解析
2. 入力データのバージョン管理
3. 不要な特徴量を取り除くための継続的な精度測定
4. パイプラインの責務の分離
5. 設定値のバージョン管理とテスタブルな設計
6. 各タスクのスケジューリングと依存管理
7. 精度のモニタリングとアラート
8. サニティチェックとしての異常値の検出

これを踏まえてワークフローを構築してみましょう。

データの取得

# daily_etl.dig
schedule:
  daily>: 00:05:00
 
+foo:
  _retry: 3
 
  +get_foo_metrics1:
    sh>: SESSION_TIME=${session_time} python elasticsearch_client.py foo metrics1 1day > foo_metrics1_daily.csv
  +vaidate_foo_metrics1:
    sh>: bash data_validator.sh foo_metrics1_daily.csv
  +commit_foo_metrics1:
    sh>: bash commit.sh foo_metrics1_daily.csv
  +get_foo_metrics2:
    sh>: SESSION_TIME=${session_time} python elasticsearch_client.py foo metrics2 1day > foo_metrics2_daily.csv
  +vaidate_foo_metrics2:
    sh>: bash data_validator.sh foo_metrics2_daily.csv
  +commit_foo_metrics2:
    sh>: bash commit.sh foo_metrics2_daily.csv

まず最初に作ったのは入力データにまつわるワークフローです。 ここではデータの取得・静的解析・バージョン管理を逐次的に行います。 各プログラムの内容についてはドメインによって大きく変わるので、以下のように定義します。

• elasticsearch_client.py.py: 引数に与えられたクラスタ名、メトリクス名、期間を元にElasticSearchからメトリクスを取得する
• data_validator.sh: 引数に与えられたデータの静的解析を行う
• commit.sh: 引数に与えられたデータをオブジェクトストレージに保存してバージョン管理を行う

Digdagでは session_time 変数としてワークフローが開始された時間を保持しているため、逐次的に処理をした場合でも各ジョブに同一の時間を与えることができます。 また、 _retry パラメータで任意の回数リトライさせることも可能です。

Digdagがワークフロー間の依存関係を持てることから、ここでは学習・予測とはワークフローを分けることで再利用性を高めています。

学習

# daily_train.dig
schedule:
  daily>: 00:10:00
 
+wait_daily_etl:
  require>: daily_etl
 
+foo:
  +train:
    _parallel: true
 
    +train_foo_metrics1:
      sh>: python train.py foo metrics1 foo_metrics1_daily.csv
    +train_foo_metrics2:
      sh>: python train.py foo metrics2 foo_metrics2_daily.csv
 
  +cross_validation:
    _pararell: true
 
    +cross_validation_foo_metrics1:
      sh>: python cross_validation.py foo metrics1
    +cross_validation_foor_metrics2:
      sh>: python cross_validation.py foo metrics2
 
_error:
  sh>: bash notify.sh

学習を行うにはまず事前にデータを取得している必要があります。 Digdagでは require オペレータでワークフロー間に依存関係を持たせる事ができるので、これによってデータの取得を待ってから学習を行う train.py を実行します。 _pararell パラメータを指定すると以下のジョブを並列に実行することができます。

また、学習を行う際には過学習に気をつけなければなりません。 そこで交差検証をして過学習が起きていた際にexit status 1を返す cross_validation.py を用意して、エラー時に発火する _error パラメータを利用して精度のモニタリングを行います。

fine-tuning

転移学習を適用して日次の学習結果を利用した週次のfine-tuningを行いたいといった複雑な依存関係も定義することができます。

# fine_tuning.dig
schedule:
  weekly>: Mon,01:05:00
 
+wait_daily_train:
  loop>: 7
  _do:
    require>: daily_train
    session_time: ${moment(last_session_time).add(i, 'day').format()}
 
+foo:
  +train:
    _parallel: true
 
    +fine_tuning_foo_metrics1:
      sh>: python fine_tuning.py foo metrics1
    +fine_tuning_foo_metrics2:
      sh>: python fine_tuning.py foo metrics2

loop オペレータで require による待ち受けを繰り返しで行うことによって日時の学習結果が出揃ったら実行という依存関係を実現しています。

予測

もちろんdaily, weekly以外にも柔軟なスケジュール定義が可能です。 予測といった頻繁に実行したいワークフローは以下のように定義します。

# predict.dig
schedule:
  minutes_interval>: 5
 
+foo:
  +predict:
    _parallel: true
 
    +predict_foo_metrics1:
      sh>: python predict.py foo metrics1
    +predict_foo_metrics2:
      sh>: python predict.py foo metrics2
 
_error:
  sh>: bash notify.sh

先ほどと同様に異常値が検出されたらexit status 1を返すようにしてアラートすることでモデルが正常に予測できているかどうかをモニタリングすることが可能です。

予測とは別のワークフローでPrediction Biasを記録してモニタリングするといったアプローチも有効になります。 その際は require オペレータで predict.dig を依存とするとうまく動くでしょう。

MItamaeによるワークフロー定義ファイル生成の自動化

ワークフローエンジンはDigdag以外にも多く存在しますが、Digdagの特徴はDSLでプログラマブルにワークフローを定義できる点にあると思います。 そこで今回はワンバイナリで動いてmrubyでプロビジョニングを記述できるMItamaeを利用して、プロビジョニング時にワークフロー定義ファイルも自動生成するようにしました。

yaml = YAML.load_file(File.join(__dir__, '../config.yml'))
file '/opt/ml/predict.dig' do
  jobs = yaml.each_with_object('') do |(key, value), string|
    value.keys.each do |metrics_name|
      string << <<-EOS
  +#{key}:
    +predict:
      _parallel: true
 
      +predict_#{key}_#{metrics_name}:
        sh>: python predict.py #{key} #{metrics_name}
EOS
    end
  end
  content <<-EOS
schedule:
  minutes_interval>: 5
 
#{jobs}
 
_error:
  sh>: bash notify.sh
EOS
end

これはMItamae(Itamae)の表現力とDigdagのDSLがあってこそで、事前に機械学習に用いる設定値とともにワークフロー生成に必要な情報を記述した以下のような config.yml を利用してワークフローを自動生成します。 そのため異なる設定値もこの config.yml に設定を追加するだけでデプロイすることが可能となります。

# config.yml
foo:
  metrics1:
    fine_tuning:
      - 'weekly'
    lstm:
      unit: 10
      window_size: 10
    metrics:
      elasticseaerch:
        index: 'foo_metrics1'
  metrics2:
    lstm:
      unit: 6
      window_size: 6
    metrics:
      elasticseaerch:
        index: 'foo_metrics2'

ワンバイナリで動くプロビジョニングツールということもあり、別の環境で十分なテストをしてからプロダクションに入れるといったような事も容易になるでしょう。 これによってワークフローの構成管理をコード化してレビューフローに乗せるだけでなく、設定値についてもそのプラクティスを適用し、同時にテストしやすい仕組みも実現することができました。

最後に

疑似コードだらけになってしまいましたが、論文中にもあるように機械学習を用いたシステムを安定運用するためにはエンジニアとリサーチャーの歩み寄りが重要であると考えています。 実際にこうしたワークフローエンジンを導入することでエンジニアリングのノウハウを機械学習を用いたシステムにも適用することができました。 今回は小規模なシステムのため一人で構築してしまいましたが、論文中にはエンジニアとリサーチャーが同じチームで働いているという記述もあり、大規模なシステムになるにつれその重要性は高まるでしょう。

機械学習を用いたシステムはまだまだ運用実績が少ないですが、技術的負債の面にもしっかりと目を向けて安定運用を目指していきたいです。 ご利用は計画的に。