Python: MLflow Models の Custom Python Models でデータを Pickle 以外に永続化する

以前、このブログでは MLflow Models の使い方について以下のようなエントリを書いた。この中では、Custom Python Models を作るときに、データを Python の Pickle 形式のファイルとして永続化していた。今回は、それ以外のファイルにデータを永続化する方法について書いてみる。

blog.amedama.jp

使った環境は次のとおり。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H2
$ python -V
Python 3.8.5
$ pip list | grep -i mlflow
mlflow                    1.11.0

下準備

まずは、下準備として MLflow をインストールしておく。

$ pip install mlflow

Pickle 形式のファイルにモデルを永続化する

まずは、おさらいとして Custom Python Models を作るときに、モデルのデータを Pickle 形式のファイルに永続化する方法から。

以下にサンプルコードを示す。この中では、定数を入力に加えるだけのモデルとして AddN というクラスを定義している。このクラスは mlflow.pyfunc.PythonModel を継承しているため、mlflow.pyfunc.save_model() と mlflow.pyfunc.load_model() を使ってファイルに読み書きできる。サンプルコードでは、実際に定数として 5 を加える設定にしたインスタンスをファイルに永続化している。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from mlflow import pyfunc


class AddN(pyfunc.PythonModel):
    """指定した数を入力に加えるモデル"""

    def __init__(self, n):
        self.n = n

    def predict(self, context, model_input):
        return model_input.apply(lambda column: column + self.n)


def main():
    # Python の Pickle ファイルとしてモデルを永続化する
    model_path = 'add-n-pickle'
    add5_model = AddN(5)
    pyfunc.save_model(path=model_path,
                      python_model=add5_model)



if __name__ == '__main__':
    main()

上記のポイントは、mlflow.pyfunc.save_model() の python_model にインスタンスを指定しているだけというところ。この場合、永続化されるのはインスタンスを Pickle 形式で直列化したファイルになる。

上記を実行してみよう。

$ python saveaddnpkl.py

実行すると、以下のように MLflow Models のフォーマットに沿ってディレクトリができる。この中で python_model.pkl が前述した AddN クラスのインスタンスを表す Pickle 形式のファイルになる。

$ ls add-n-pickle 
MLmodel         conda.yaml      python_model.pkl
$ python -m pickle add-n-pickle/python_model.pkl
<__main__.AddN object at 0x1048cdd90>

上記を Python から読み込んで使うサンプルコードも次のように用意した。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pandas as pd
from mlflow import pyfunc


def main():
    model_path = 'add-n-pickle'
    loaded_model = pyfunc.load_model(model_path)

    x = pd.DataFrame(list(range(10, 21)), columns=['n'])
    y = loaded_model.predict(x)

    print(f'Input: {x.n.values}')
    print(f'Output: {y.n.values}')


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python loadaddnpkl.py
Input: [10 11 12 13 14 15 16 17 18 19 20]
Output: [15 16 17 18 19 20 21 22 23 24 25]

ちゃんと、読み込んだモデルが入力データに定数として 5 を加えていることが確認できる。

テキストファイルにインスタンスのアトリビュートを永続化してみる

それでは、続いて Pickle 以外のフォーマットのファイルも使って永続化するパターンを扱う。これは、たとえば永続化したいモデルが Pickle 以外のフォーマットでファイルに読み書きする API がある場合などに使い勝手が良い。実際に MLflow Models と XGBoost や LightGBM のインテグレーションは、フレームワークが提供する永続化用 API を流用して書かれている。

以下にサンプルコードを示す。今回は、入力に加算する定数をテキストファイルとして、AddN クラスのインスタンスとは別に永続化するやり方を取った。また、先ほどはファイルへの書き込みと読み込みを別の Python モジュールに分けたのに対して、これはひとつで完結させている。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import tempfile
import os

import pandas as pd
from mlflow import pyfunc


class AddN(pyfunc.PythonModel):
    """指定した数を入力に加えるモデル"""

    def __init__(self, n=None):
        self.n = n

    def load_context(self, context: pyfunc.PythonModelContext):
        """インスタンスの状態を復元するのに使われるコールバック"""
        # アーティファクトのパスを取り出す
        artifact_path = context.artifacts['n_file']
        with open(artifact_path, mode='r') as fp:
            # ファイルからモデルのパラメータを復元する
            self.n = int(fp.read())

    def predict(self,
                context: pyfunc.PythonModelContext,
                model_input: pd.DataFrame):
        return model_input.apply(lambda column: column + self.n)

    def __repr__(self):
        """インスタンスの文字列表現を得るときに呼ばれる特殊メソッド"""
        return f'<AddN n:{self.n}>'


def save_model(n: int, path: str):
    """モデルをアーティファクトに永続化するためのユーティリティ関数"""
    # モデルが動作するのに必要なパラメータをテキストファイルなどにアーティファクトとして書き出す
    filename = 'n.txt'
    with tempfile.TemporaryDirectory() as d:
        # 一時ディレクトリを用意して、そこにファイルを作る
        artifact_path = os.path.join(d, filename)
        with open(artifact_path, mode='w') as fp:
            fp.write(str(n))
        # 上記で作ったファイルをアーティファクトとして記録する
        # NOTE: path 以下にファイルがコピーされる
        artifacts = {
            'n_file': artifact_path,
        }
        pyfunc.save_model(path=path,
                          # このインスタンスに load_context() メソッド経由でパラメータが読み込まれる
                          python_model=AddN(),
                          artifacts=artifacts,
                          )


def main():
    # モデルをアーティファクトに永続化する
    model_path = 'add-n-artifacts'
    save_model(5, model_path)

    # モデルをアーティファクトから復元する
    loaded_model = pyfunc.load_model(model_path)

    # 動作を確認する
    x = pd.DataFrame(list(range(10, 21)), columns=['n'])
    y = loaded_model.predict(x)

    print(f'Input: {x.n.values}')
    print(f'Output: {y.n.values}')


if __name__ == '__main__':
    main()

上記にはポイントがいくつかある。まず、書き込みに関しては mlflow.pyfunc.save_model() を呼ぶ際に artifacts というオプションを指定している。これには、モデルのデータなどを記録したファイルへのパスを Python の辞書として渡す。ちなみに、ここに指定するパスは実際には Artifact URI なので、リモートのストレージにあっても構わない。この点は、おそらく MLflow Tracking と組み合わせて使うことを想定しているのだと思う。もちろん、ファイルはあらかじめ、そのパス (繰り返しになるけど、実際には Artifact URI) に存在している必要がある。ここに指定したファイルは、MLflow Models が作成するディレクトリへとコピーされる。そして、読み込みに関しては mlflow.pyfunc.PythonModel を継承したクラスに load_context() というメソッドを実装する。このメソッドは、インスタンスを Pickle 形式のファイルから読み込んだ後に呼ばれるようだ。つまり、Pickle 以外のファイルにデータを保存しているときに、モデルの状態をそれで更新するときに使うことができる。

さて、前置きが長くなったけど実際に上記のサンプルコードを実行してみよう。

$ python addnartifacts.py
Input: [10 11 12 13 14 15 16 17 18 19 20]
Output: [15 16 17 18 19 20 21 22 23 24 25]

ちゃんと想定どおりの入出力になっている。

MLflow Models が作成したディレクトリを確認してみよう。すると、artifacts というサブディレクトリがあることに気づく。

$ ls add-n-artifacts 
MLmodel         artifacts       conda.yaml      python_model.pkl

中を見ると、インスタンスのアトリビュートがテキストファイルとして書き込まれている。

$ cat add-n-artifacts/artifacts/n.txt   
5

それ以外には、アトリビュートが None の状態の AddN クラスのインスタンスが Pickle 形式で永続化されていることがわかる。

$ python -m pickle add-n-artifacts/python_model.pkl
<AddN n:None>

いじょう。

Linuxで動かしながら学ぶTCP/IPネットワーク入門

作者:もみじあめ
発売日: 2020/02/29
メディア: Kindle版

CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: MLflow Models の Custom Python Models でデータを Pickle 以外に永続化する

もくじ

下準備

Pickle 形式のファイルにモデルを永続化する

テキストファイルにインスタンスのアトリビュートを永続化してみる