【Techの道も一歩から】第44回「Prefectでパイプラインを構築してみる」

f:id:kanjirz50:20190104142720j:plain

こんにちは。 DSOC R&D グループの高橋寛治です。

本記事は Sansan Advent Calendar 2021 - Adventar の初日の記事です*1。

Pythonによるパイプラインツールについて、Prefectというフレームワークが良さそうという話を小耳に挟みました。

第42回の記事と同様にMNISTを題材にパイプラインを作ってみます。

buildersbox.corp-sansan.com

Prefect

Prefectは、Pythonで書かれたオープンソースのワークフロー管理システムです。

特徴は、タスクを定義し、タスクを利用したフローを書くというシンプルな構成となっていることです。タスクはフローを構成する処理を表します。タスクとフローを定義する箇所以外は、Prefectが吸収してくれます。

説明よりも実際のコードを見た方が早いと思いますので、MNISTを題材に見ていきます。

パイプラインを構築する

MNISTデータセットを取得し、前処理、学習、評価を行う単純なパイプラインを構築します。

タスクは、クラスもしくは関数で定義します。ここでは、関数によるタスクを例に挙げます。

from prefect import task, Flow
from sklearn.datasets import load_digits


@task
def load_digits_data():
    return load_digits()


def build_mnist_experiment_flow():
    with Flow("mnist-experiment-flow") as flow:
        digits_data = load_digits_data()

    return flow


flow = build_mnist_experiment_flow()
flow.run()

タスクを実装する関数にtaskデコレータを適用します。

フローについては、with文の中で定義します。上記の例だと、build_mnist_experiment_flow関数で何をするかを記述します。

次にパラメータです。安直ですが、学習とテストの分け方の乱数の種をパラメータで与えることを考えます。

from prefect import Parameter

# 中略

@task
def preprocesing_task(digits_data, random_state):
    X_train, X_test, y_train, y_test = train_test_split(
        digits_data.images.reshape((len(digits_data.images), -1)),
        digits_data.target,
        test_size=0.5,
        shuffle=True,
        random_state=random_state
    )

    dataset = {
        "X_train": X_train,
        "X_test": X_test,
        "y_train": y_train,
        "y_test": y_test
    }

    return dataset


def build_mnist_experiment_flow():
    with Flow("mnist-experiment-flow") as flow:
        random_state = Parameter("random_state", default=42)

        digits_data = load_digits_data()
        dataset = preprocesing_task(digits_data, random_state)

    return flow


flow = build_mnist_experiment_flow()
flow.run(random_state=42)