Kubeflow Pipelines (kfp)

Kubeflow Pipelines と W&B を統合して、ML パイプラインを構成する各コンポーネントにまたがる Experiments を管理し、メトリクスを可視化します。

Kubeflow Pipelines (kfp) は、Dockerコンテナーをベースとした、移植性とスケーラビリティに優れた機械学習 (ML) ワークフローを構築・デプロイするためのプラットフォームです。

このインテグレーションを使用すると、kfp の Python 関数コンポーネントにデコレータを適用して、パラメーターとアーティファクトを自動的に W&B にログできます。

この機能は wandb==0.12.11 で利用可能になり、kfp<2.0.0 が必要です

登録してAPIキーを作成する

APIキーは、お使いのマシンをW&Bに対して認証するために使用します。APIキーはユーザープロフィールから発行できます。

より手早く行うには、User Settings に直接アクセスしてAPIキーを作成してください。新しく作成したAPIキーはすぐにコピーし、パスワードマネージャーなどの安全な場所に保存してください。

1. 右上にあるユーザープロフィールアイコンをクリックします。
2. User Settings を選択し、API Keys セクションまでスクロールします。

wandb ライブラリをインストールしてログインする

wandb ライブラリをローカルにインストールしてログインするには、次の手順に従います。

コマンドライン

1. WANDB_API_KEY 環境変数 に APIキーを設定します。

   export WANDB_API_KEY=<your_api_key>
   

2. wandb ライブラリをインストールし、ログインします。

   pip install wandb
   
   wandb login
   

Python

pip install wandb

import wandb
wandb.login()

Python notebook

!pip install wandb

import wandb
wandb.login()

コンポーネントをデコレートする

@wandb_log デコレータを追加し、通常どおりコンポーネントを作成します。これにより、パイプラインを実行するたびに、入力/出力のパラメーターとアーティファクトが自動的に W&B にログされます。

from kfp import components
from wandb.integration.kfp import wandb_log


@wandb_log
def add(a: float, b: float) -> float:
    return a + b


add = components.create_component_from_func(add)

コンテナーに環境変数を渡す

環境変数をコンテナーに明示的に渡す必要がある場合があります。双方向リンクを有効にするには、環境変数 WANDB_KUBEFLOW_URL に Kubeflow Pipelines インスタンスのベース URL も設定してください。たとえば、https://kubeflow.mysite.com です。

import os
from kubernetes.client.models import V1EnvVar


def add_wandb_env_variables(op):
    env = {
        "WANDB_API_KEY": os.getenv("WANDB_API_KEY"),
        "WANDB_BASE_URL": os.getenv("WANDB_BASE_URL"),
    }

    for name, value in env.items():
        op = op.add_env_variable(V1EnvVar(name, value))
    return op


@dsl.pipeline(name="example-pipeline")
def example_pipeline(param1: str, param2: int):
    conf = dsl.get_pipeline_conf()
    conf.add_op_transformer(add_wandb_env_variables)

プログラムでデータにアクセスする

Kubeflow Pipelines UI から

W&B にログされている Kubeflow Pipelines UI 内の任意の Run をクリックします。

• Input/Output タブと ML Metadata タブで、入力と出力の詳細を確認できます。
• Visualizations タブから W&B Web アプリを表示できます。

web app UIから

web app UI には、Kubeflow Pipelines の Visualizations タブと同じ内容が、より広い表示領域で表示されます。web app UI の詳細はこちら。

Public API を使用する (プログラムからアクセスする場合)

• プログラムからアクセスするには、Public APIを参照してください。

Kubeflow Pipelines と W&B の概念対応

以下は、Kubeflow Pipelines の概念と W&B の対応関係です

Kubeflow Pipelines	W&B	W&B での場所
入力スカラー	config	Overview タブ
出力スカラー	summary	Overview タブ
入力アーティファクト	入力アーティファクト	Artifacts タブ
出力アーティファクト	出力アーティファクト	Artifacts タブ

より細かなログ制御

ログをより細かく制御したい場合は、コンポーネント内で wandb.log と wandb.log_artifact を適宜呼び出せます。

wandb.log_artifacts を明示的に呼び出す

以下の例では、モデルをトレーニングしています。@wandb_log デコレータは、関連する入力と出力を自動的にトラッキングします。トレーニングプロセスもログしたい場合は、次のようにそのログを明示的に追加できます。

@wandb_log
def train_model(
    train_dataloader_path: components.InputPath("dataloader"),
    test_dataloader_path: components.InputPath("dataloader"),
    model_path: components.OutputPath("pytorch_model"),
):
    with wandb.init() as run:
        ...
        for epoch in epochs:
            for batch_idx, (data, target) in enumerate(train_dataloader):
                ...
                if batch_idx % log_interval == 0:
                    run.log(
                        {"epoch": epoch, "step": batch_idx * len(data), "loss": loss.item()}
                    )
            ...
            run.log_artifact(model_artifact)

暗黙的なwandbインテグレーションを使用する場合

サポート対象のフレームワークインテグレーション を使用している場合は、コールバックを直接渡すこともできます。

@wandb_log
def train_model(
    train_dataloader_path: components.InputPath("dataloader"),
    test_dataloader_path: components.InputPath("dataloader"),
    model_path: components.OutputPath("pytorch_model"),
):
    from pytorch_lightning.loggers import WandbLogger
    from pytorch_lightning import Trainer

    trainer = Trainer(logger=WandbLogger())
    ...  # トレーニングを実行する