🎃

続・TFMAで始めるModelAgnosticなモデル分析

2020/12/23に公開

TensorFlow

tfx

tech

初めに

TensorFlow Advent Calendar 23日目の記事です
以前BigQuery Advent Calenderの記事として書いたTFMAで始めるBigQuery MLもモデル分析の続編です。

この記事の中では一度ローカルにpandas.DataFrameとしてデータをダウンロードをしてきて、そこからTFMAにデータを渡していました。
この方法だとDataFrameを可視化したりすることでデバッグが容易にできるので、TFMAのようなまだまだ絶賛開発中のライブラリを触る分には助かることが多いです。

一方で、DataFrameに収まらない規模のデータを扱おうとしたときには、高性能のシングルマシンを用意して物理で殴るみたいなことをする必要があります。
TFMAは内部でApache Beamを用いてデータの処理をしてモデル分析に必要な評価値の計算をしています。なので、Google Cloud DataFlowを使うことでBigQueryに保存されている予測結果について、データの規模によらずシームレスにTFMAを実行できるのではと挑戦しました。

利用したデータセットはbigquery-public-data.ml_datasets.credit_card_defaultです。
選定理由は、Result Slicing with Cloud AutoML TablesというTFMA+AutoML Tablesのチュートリアルで利用されていたので、同じようなことをしたい場合に利用できるだろうと考えたためです。

手順

前回の手順3.で作成したtfma.SlicingSpecやtfma.EvalConfigはこちらの手順でも同様に必要です。また、GCSに分析結果の成果物を書き込むことを想定していますので、バケットの作成も必要です。

以下手順です。それぞれの手順でサンプルになるコードも一緒に記載しておきます

*注意: 必要最低限の記載しかしていないのでコピペすれば動くというものではありません
後日colabで動くノートブックを公開します。

default_writersを用意する
BigQueryのデータをのpyArrow.RecordBatchに変換する
tfma.ExtractEvaluateAndWriteResultsへ渡して、結果を書き込む
モデル分析の成果物を読み込んで、可視化する

1. default_writerを用意する

tfma.default_writersは手順3.で利用するたモデルの分析結果を書き込むためのWriterです。
今回はModel非依存のModelAgnosticなものなので、結果書き込みパスとtfma.EvalConfigを引数に渡して作成します。

writers = tfma.default_writers(
  output_path,
  eval_config=eval_config
)

2. BigQueryのデータをのArrow.RecordBatchに変換する

ここからはApache Beam Pipelineの作業になります。
TFMAにデータを渡すために、BigQuery からデータを抽出してApacheArrowのRecordBatchに変換する処理を行います。
データを保持するオブジェクトは Dict -> pandas.DataFrame -> pyarrow.RecordBatch(ListArray)のように変換していきます。
ここで、一度DataFrameに変換しているのは、TFXの便利関数群のtfx-bslのtable_utilsを利用したいためです。

変換するコードは以下のようになります。

def ConvertRecordBatch(data):
  # Beamのプログラミングパラダイムでどこでインポートするのがいいのか迷ったので
  # ここでインポート
  from tfx_bsl.arrow import table_util 
  import pandas as pd
  return table_util.CanonicalizeRecordBatch(table_util.DataFrameToRecordBatch(pd.DataFrame.from_dict(data={k:[v] for k, v in data.items()})))
  
with beam.Pipeline() as pipeline:
  _ = (
      pipeline
       | "QueryTable" >> beam.io.gcp.bigquery.ReadFromBigQuery(
           gcs_location=output_path + "/bigquery/",
           use_standard_sql=True,
           query = query,
           )
       | "ConvertArrow" >> beam.Map(ConvertRecordBatch)
  )

3. tfma.ExtractEvaluateAndWriteResultsへ渡して、結果を書き込む

2のパイプラインの続きです。1で作成したtfma.default_writer、tfma.EvalConfig、結果の書き込み先のoutput_pathを引数に渡して、この手順は完了です！

def ConvertRecordBatch(data):
  # Beamのプログラミングパラダイムでどこでインポートするのがいいのか迷ったので
  # ここでインポート
  from tfx_bsl.arrow import table_util 
  import pandas as pd
  return table_util.CanonicalizeRecordBatch(table_util.DataFrameToRecordBatch(pd.DataFrame.from_dict(data={k:[v] for k, v in data.items()})))

output_path = "gs://bucket/path/to/result"

writers = default_writers(
    output_path,
    eval_config=eval_config
)

with beam.Pipeline() as pipeline:
  _ = (
      pipeline
       | "QueryTable" >> beam.io.gcp.bigquery.ReadFromBigQuery(
           gcs_location=output_path + "/bigquery/",
           use_standard_sql=True,
           query = query,
           )
       | "ConvertArrow" >> beam.Map(ConvertRecordBatch)
       | "ExtractEvaluateAndWriteResults" >> ExtractEvaluateAndWriteResults(
           writers=writers,
           eval_config=eval_config,
           output_path=output_path)
  )

4. モデル分析の成果物を読み込んで、可視化する

3まででめんどくさい準備は終わりました。あとは読み込んで可視化するだけです。

result = tfma.load_eval_result(output_path)
tfma.view.render_slicing_metrics(result)

躓いた点

このやり方でいくつか躓いたり引っかかったりしたところを共有しておこうと思います。

ExtractEvaluateAndWriteResultsを実行時に、ラベルや予測値がStringになってるとエラーになる
tf.train.Example形式のpCollectionをExtractEvaluateAndWriteResultsに渡しても、エラーになる(?)

2.についてはちょっと調査が足りてないのですが、tf.train.Example形式のデータを渡してもエラーになりました。ExtractEvaluateAndWriteResultにはいろいろな引数が渡せるのですが、今回はeval_config, output_path,writersしか利用していません。

例えば、Result Slicing With Cloud AutoML Tablesを見てみると、tfma.model_agnostic_evalを利用してEvalSharedmodelとextractorsというものを作成してます。この場合には、tf.train.Exampleのデータを渡しても動くのではないかなと思います。

楽するために、TFMAで始めるBigQuery MLもモデル分析のコードや設定の仕方をそのまま使い回したかったので、このやり方を採用しました。

未調査な点

NULL値の扱い。
table_utilsのなかで、null値にマスクをするような処理が見られましたが、あまり追っていないので最終的にモデルを評価するときにNull値をどうしてるのかまでは確認できていません。
躓いた点2で記載した、tf.train.Example形式を利用したModelAgnosticなTFMAの動かし方
BigQuery -> TFDVをBeam(Dataflow)で動かしたときにはtf.train.Exampleでうまくいったような記憶があるので、TFMAでも楽に使うことができるのではと期待しています。

(3. その他思い浮かばなかったこと)

終わりに: TFXの関連ライブラリを触る時のコツ

TFX関連の記事があまりにも少なく、尻込みしている人も多いのではと思います。大雑把にTFXというものが何なのかを知るには、Towards ML Engineering: A Berief History Of TensorFlow Extended(TFX)を読むのがいいと思います。
日本語の情報であれば、昨年のアドベントカレンダーの記事のTFXとは何だったのか、現状どうなっているのかが簡潔にまとまっていて良いです。

TFXについてはPRを投げるのはともかく、利用していく上では高レベルAPIが用意されているので比較的楽にある程度安定して使えると思います。

ここでは、心構えとか、どう実装するのか迷った時にどう調べたりすればいいのか、というのを自戒を込めて書いておこうと思います。

心構えとしては、まずサンプルは動かない。サンプル以上のことをしようとすると動かないといううのを念頭に置いておきましょう。実際に動かないことはないと思いますが、最初からよしなに動くものを期待していくとエラーばっかり出てしんどくなります。
そもそも0.25みたいなバージョンのライブラリに動くものを期待しても仕方ないです。
諦めて慣れましょう。その内、思った通りに動かないことに安心感を覚えるようになります。

現行バージョン実装の仕方に悩んだ時はTFXの各Exectutorに目を通しましょう。
TFXはざっくり言えば、I/Oを定義したComponentと実際の内部処理のExecutorに分かれています。
Executorの実装を見れば、ある程度の大雑把な使い方は理解できると思います。
TFMAの例で言えば, evaluatorまたは、model_validatorをみると、どういう実装が期待されているのかとかがわかると思います。

ちょっと改良して自分たちのシステムに組み込みたいという場合には、高レベルAPIの処理を追っていくのも手です。今回だと、tfma.run_model_analysisやtfma.analyze_raw_dataを追っていくのがいいと思います。

余力があれば、protoファイルを確認しておくと、データの流れをスムーズに把握できていいと思います。TFX関連で使われるprotoファイルはtensorflow_metadataにまとまっています。各ライブラリ専用のものは個別のリポジトリにあります。

例えば、統計情報やスキーマ情報なんかはTFDVやTensorFlow-datasetsに使われたりするので、tensorflow_metadataにまとめていて各ライブラリのprotoファイルでimportして利用していますが、tfma.EvalConfigはTFMA専用のものなので、tensorflow_model_anlaysis/protoにまとまっています。

TFXは関連ライブラリもパイプラインも実装を読んでいくといろいろな発見があるので、単純に読み物としても面白いと思います。
特に、システムとしてMLを組み込みたい勢にはハマるのではと思っています。

P.S.
(高レベルAPIを利用しないで使ってみようとした瞬間に把握することが多くなり、何もわかんない状態になったのに加えて2日前くらいから空き時間で触り始めたので大変でした😇)