kedro触ってみた

PyTorchで学習周りのコードを書いてたのですが、今後も検証を勧めていく上で、実験結果の記録やパラメータ変更や複数モデルの比較などをするために、何らかの学習周りのコードを扱うフレームワークを利用しようと思っていました。

PyTorchやTensorFlowなどの外側の部分について、簡単にまとめつつ、kedroを触ってみたよという記事になります。

機械学習のライブラリについて

モデルの構築にはPyTorchやTensorFlowが挙げられますが、それらより外側の物事を扱う周辺ライブラリはデファクトスタンダードが決まっていないように思います。

参考：https://qiita.com/fam_taro/items/c32e0a21cec5704d9a92

上記の記事で挙げられるようなものは、学習周りのコードを定めるためのフレームワークです。これらのタイプのライブラリに、ちょうどいい分類の名前が無いので、「学習周りのコードを定めるためのフレームワーク」としか言えないのが難しいところです。
IgniteもPyTorch Lightningも、READMEにはハイレベルのAPIであるとか、オーガナイズするラッパーですとしか書かれていません。

よく考えてみると、Rails, Laravel, DjangoあるいはFlaskといったwebアプリケーションフレームワークという呼称も、知っている人からすればどういうものなのかは想像しやすいですが、「webアプリケーションフレームワーク」という名前はそれらの実体を指すには不明瞭であるように思えます。

IgniteかPyTorch Lightningかをとりあえず触ってみようかとも思ったのですが、今回は

• 定義した学習のコードを、他のデータですぐに学習を開始できる
• 実験ごとにjupyter notebook書き捨てせず、後から見た人に状況を説明しやすくしたい

以上のようなことを踏まえ、パイプライン構築のためのライブラリを触ってみることにしました。
（あとからパイプライン内部で、上記のようなライブラリを利用することもできるはずです。）

パイプライン構築のためのライブラリ

パイプラインのためのライブラリはこちらの記事によくまとまっています。
PythonのPipelineパッケージ比較：Airflow, Luigi, Gokart, Metaflow, Kedro, PipelineX

なお、こちらの記事の筆者はPipelineXの作者であることはご注意下さい。ただそのことを割り引いて見ても、実際私もkedroを触ってみた後に、いくつか気になるところがあったのでそれらを解決してくれるラッパーが欲しくも感じました。いずれPipelineXの方も確認してみたいと思います。

今回は、以下のような点から、kedroを使ってみました。

• 拡張性が高そうである
• node, pipelineという概念が単純
• kedro-vizという視覚化ツールがある
• スターが多くてユーザが多そう
• ドキュメントが多い
• ソースコードがテストを除いて１万行未満

また、最悪の場合として、コードが肥大した上でkedroをやめる場合でも、nodeとして定義した関数を結合するだけなので他のライブラリに乗り換える労力も少ないのではないかと考えました。

kedro

https://github.com/quantumblacklabs/kedro

Kedro is an open-source Python framework that applies software engineering best-practice to data and machine-learning pipelines.

ドキュメント
kedroとは何なのかについては、公式ドキュメントやその他ググって見つけられる記事におまかせします。

既存のコードをkedroに置き換える

既にjuypter notebookやpythonスクリプトで定義したコードを、ちょうどよい粒度で、ドキュメントに従ってnodeとして分けていきます。

色々と省略していますが、Pipelineを作る一例です。

from kedro.pipeline import Pipeline, node
from .get_optimizers import get_optimizers
from .fit import fit

def create_pipeline():
    return Pipeline(
        [
            node(
                get_optimizers,
                inputs=["pretrained_model@load", "total_steps"],
                outputs="optimizers",  # tuple of optimizer and scheduler
            ),
            node(
                fit,
                [
                    "pretrained_model@load",
                    "optimizers",
                    "example_train",
                    "example_valid",
                    "params:epochs",
                ],
                "model",
            ),
        ],
    )

Nodeとしてまとめたい内容を関数として渡せばいいので、既存のコードを別ファイルとして書いておき、importして渡すだけです。

Nodeという概念は、関数とほぼ同義ですが、入出力に名前をつけられるという点が便利と言えます。

実際のコードは省略しますが、自分が作ったものをkedro-vizで図にすると以下のようになりました。

青い文字と枠は私が追記したものです。

この図だけですべてを分かってもらうのは難しいかも知れないですが、補足する図として利用するにはとても有用に思えます。直接伝えたりコードを読んで理解したりする上でこの図の有無で結構違うのではないかと思います。

遭遇した注意点

パイプラインのNodeとNodeの間で受け渡されるデータは、kedroのDataSetというクラスで扱われます。
特にcatalg.ymlで設定をしないとMemoryDataSetが使われ、メモリと言う名前の通り、ファイルシステム等で永続化はされません。
このときの受け渡し、すなわちメモリ上での挙動が、デフォルトだとdeep copyです。CSVで表現されるようなテーブル上の単なるデータならコピーで問題ないのですが、機械学習で用いられる何らかのpythonオブジェクトを受け渡す際は、当然これはコピーではまずいです。
MemoryDataSetはこのコピーの挙動を指定できるので、assignとすることで、pythonオブジェクトがコピーされずそのまま受け渡せるようになります。

optimizers:
  type: MemoryDataSet
  copy_mode: assign

これ自体はドキュメントにも書かれているのですが、PySparkとの連携に関するドキュメントの中に書かれており、この挙動に気づくのにしばらく時間がかかりました。
（コピーされたオブジェクトが渡されていてエラーが発生しないので、学習を回しても精度が動かず原因特定に悩まされました）

感じた不満点

• ドキュメントの質
  ドキュメントは確かに多いのですが、getting startedだけを読んですぐに小規模な実行が始められるという状態ではなく、結果として隅々まで読むことになりました。（もちろん、ドキュメントが無いよりかは十分良い状態ではあります） 特に、A “Hello World” exampleからの飛躍が大きく感じます。
  ここのパート自体は単純なのですが、次のパートからkedro newで作られるプロジェクトフォルダでの作業になりますが、もう少しミニマルな構成が欲しくも感じました
• データのバージョニングのまとめ方
  まだドキュメントを読んだ程度の知識ですが、kedroにはデータのバージョニング自体の機能はあるのですが、dataフォルダへの保存の仕方に違和感を覚えました。
  バージョンごとにフォルダが分かれるのではなく、個別のデータに対してフォルダが作られ、そのしたにバージョンidで異なるデータが保存されるようになっています。これはポータビリティの観点や、途中実行のしやすさから気持ち悪く感じます。
  しかも、この機能はDataSetクラスを継承するときに定義する形で実装されてるため、今後この挙動が変更される可能性は低いように思います。

まとめ

普通に使い始めるまでに、色々とドキュメントを読み漁ったり、上記した注意点に気づかずに戸惑ったりしましたが、何もない状態よりかはマシかもしれないな、と感じることができました。
ただ、「１実験１notebook」という運用でゴリゴリやるのと比べて飛躍的に便利であるかというと、まだ判断ができません。
ただ、拡張性が高いのでプロジェクト固有の処理などを定義したり差し込み処理などが実装しやすいようにはなっているようなので、ちょっとした気に入らない部分があれば自分で拡張して解決もできそうです。