kedroのモデルを読み込むパイプラインにおいて、同じ型のモデルを切り替える

kedroのパイプラインで、同じ型のモデルを差し替えて利用する方法を、私が分かった＆やってみた範囲で紹介します。

例えば、BERTとELECTRAにおいて、どちらも各々の事前学習モデル（のパラメータ）があるとして文書分類タスクをファインチューニングします。
これを行うようなパイプラインを定義して、全体の構造は変えずにこの部分的なパイプラインを２つのモデルで切り替えられるようにしました。

先に結論を書いておくと、同じ型に基づいてパイプラインを定義して…というきれいな方法はやらずに、愚直にifで切り替える方法にとどめています。
（もしも、抽象クラスで定義したパイプラインに各々の実装を持ったクラスを注入する方法を見たかった人がいたら、すみません！）
途中のDataSetを繋ぐnodeを追加するのが肝であると思います。

準備

kedroの基本的なことは既知とします。公式ドキュメントを参照してください。

dataフォルダに、BERTとELECTRAの事前学習済みモデルを置いておきます。
ここのパスを、catalog.ymlでDataSetとして定義しておくと良いかと思います。

今回はこのようなパイプラインの流れを考えます。

図の四角のように３つのパイプラインに分けて考えます。

• データの用意
• モデルのロード
• finetuning

上記で言っているBERTやELECTRAで共通化したいのは、真ん中のモデルのロードにあたります。

前後の別のパイプラインは、BERTやELECTRAについては知らずに、finetuningもただ共通の型として実行したい、ということです。

注意点：動的にDataSetは切り替えられない

パイプラインについて考える前に、kedroの仕様として知っておく必要があることですが、
catalog.yml上で、動的にDataSetを切り替えることは現状できません。
参考: https://github.com/quantumblacklabs/kedro/issues/138

なので、BERTやELECTRAごとの「事前学習モデルのパス」「事前学習モデル自体」などのDataSetは、実行したい各々のパターンを書き下す必要があります。

model-bert:
  #hogehoge
model-electra:
  #hogehoge

その上で、pipelineの定義や呼び出しに置いて、可能な範囲で分かりやすい書き方をしていくことを考えましょう。

パイプラインの結合部分について

kedroのpipelineは、名前を切り替える機能があります。今回はこれを使いました。

 pipeline(p, inputs={"hoge": "hoge2"})

ただし、これはパイプラインのdatasetの名前を置き換えるものであり、aliasのようにつなぐものではないです

なので、今回のようにパイプラインの前後に処理がある場合、
先にリネームをすると、後のパイプラインと結合できなくなります。

全体のパイプラインで、modelをmodel-bertにリネームすることは可能ですが、これは、一部のパイプラインの関心事が外側に露出しているので気持ち悪いです。
外側のパイプラインでは、model-bertのようなDataSetには触らず、あくまでmodelといったDataSet名だけを扱うようにしておきたいです。

なので、renameと合わせて、DataSet名を繋ぐためだけに何も処理をせず入力をそのまま出力するnodeをモデルのロードを行うパイプラインに追加します

 p += Pipeline(Node(lambda a: a, "model-bert", "model"))

このようなノードをモデルをロードするpipeline内で追加するようにすることで、モデルをロードするパイプライン内のみでDataSet名はBERTやELECTRAを指定しつつ、全体でパイプラインを結合する際は、外側のパイプラインとはmodelを経由して結合がができます
ただし、modelDataSetは必要に応じてcatalog.ymlにも記述する必要があるのに注意してください。

また、入力側に関しても同様にmodel_path-bertをmodel_pathに繋ぐことで、内部のロードはmodel_pathは詳細を知らずに読み込むようにできます。

パイプラインの中身について

理想を言えば、BERTやELECTRAといった具体的なクラスは、外側から注入するのがいいかも知れませんが、３つ目以上の同じ型のモデルを扱うことになる可能性自体がそもそも少ないです。

パイプラインのエントリーポイントは各pipelineモジュールのフォルダのpipeline.pyのcreate_pipeline関数です。ここにモデル名を文字列で引数として渡す形で記述するのが、このようなケースでは十分と思えます。
（ただし、モデルの設計自体をしてる場合はこの限りではないかも知れません）

また、このパイプライン内の処理自体も短いので、node用のscriptをコピーして、BERT用とELECTRA用に分けて記述するのも十分かと思われます。一般的には良くないことですが、モデルの読み込みのコードは変更が少なく規模も大きくなりにくいので、コピーするほうが単純で理解しやすいかと思えます。
２つのファイルに対し同じ変更を何度かしたり、３つ目以降のモデルを読む必要になったりしたタイミングで共通化を検討すればよいかと思います。

以上を元に、以下のようにパイプラインを作ることができます。

# tf_to_torch, tf_to_torch_electraというnode用の関数をimportしているとします

def get_pipeline(fn, **kwargs):
    return Pipeline(
        [
            node(
                fn,
                ["pretrained_model_path", "num_labels"],
                "pretrained_model_tf",
            ),
        ],
        **kwargs,
    )

def create_pipeline(model: str):
    if(model == "bert"):
        fn = tf_to_torch
    elif(model == "electra"):
        fn = tf_to_torch_electra
    else:
        raise Exception(f"not supported model: {model}")

    p = get_pipeline(fn, tags=f"get_model_{model}")
    p += Pipeline([node(lambda a: a, f"model_path-{model}", "model_path")])

    p = pipeline(p, outputs={"model_tf": f"model-{model}"})
    p += Pipeline([node(lambda a: a, f"model-{model}", "model")])
    return p

補足：pythonのpathlibについて

pythonの3.4から使える、pythlibという標準パッケージがあり、pythonにおけるファイルやフォルダのパスの操作が快適に行えるパッケージです。
os.path.join等の文字列を直接操作する方法しか知らなかった人は、pathlibを使うことをオススメします。

またkedroでは、パスをDataSetとして使う場合にそのためのDataSetクラスを定義しておくと楽です。

# src/nlp_pipeline/io/path_dataset.py
from typing import Any, Dict
from kedro.io import AbstractDataSet
from pathlib import Path

class PathDataSet(AbstractDataSet):
  def __init__(self, path: str):
    self._path = Path(path)

  def _load(self) -> Path:
    return self._path

  def _save(self, _path) -> None:
    self._path = _path

  def _describe(self) -> Dict[str, Any]:
    return dict(path=self._path)

利用例

model_path-bert:
   type: nlp_pipeline.io.path_dataset.PathDataSet
   path: data/01_raw/bert-wiki-ja/