ML ワークフローでの Dataflow ML

複雑な ML ワークフローをオーケストレートするには、データの前処理と後処理のステップを含むフレームワークを作成します。モデルのトレーニングやデータの後処理を行う前に、データの前処理が必要になることがあります。

多くの場合、ML ワークフローにはパイプラインを形成する多くのステップが含まれています。ML パイプラインを構築するには、次のいずれかの方法を使用できます。

  • Apache Beam や Dataflow ランナー(TensorFlow Extended(TFX)、Kubeflow Pipelines(KFP)など)とのインテグレーションが組み込まれたオーケストレーション フレームワークを使用します。これは最も手軽なオプションです。
  • Dataflow テンプレートでカスタム コンポーネントを作成し、ML パイプラインからそのテンプレートを呼び出します。この呼び出しに Apache Beam コードが含まれています。
  • ML パイプラインで使用するカスタム コンポーネントを作成し、そのコンポーネントに Python コードを直接配置します。カスタム Apache Beam パイプラインを定義し、カスタム コンポーネント内で Dataflow ランナーを使用します。このオプションは最も複雑で、パイプラインの依存関係を管理する必要があります。

ML パイプラインを作成したら、オーケストレーターを使用してコンポーネントを連結し、エンドツーエンドの ML ワークフローを作成できます。コンポーネントをオーケストレートするには、Vertex AI Pipelines などのマネージド サービスを使用します。

ML アクセラレータを使用する

大規模なモデルでの推論など、計算負荷の高いデータ処理を伴う ML ワークフローでは、Dataflow ワーカーでアクセラレータを使用できます。Dataflow は、GPU と TPU の両方の使用をサポートしています。

GPU

処理を高速化するには、Dataflow ジョブで NVIDIA GPU を使用します。Dataflow は、T4、L4、A100、H100、V100 など、さまざまな NVIDIA GPU タイプをサポートしています。GPU を使用するには、必要な GPU ドライバとライブラリがインストールされたカスタム コンテナ イメージを使用してパイプラインを構成する必要があります。

Dataflow で GPU を使用する方法については、Dataflow での GPU のサポートをご覧ください。

TPU

Dataflow は、大規模な AI モデル向けに最適化された Google のカスタム設計の AI アクセラレータである Cloud TPU もサポートしています。TPU は、PyTorch、JAX、TensorFlow などのフレームワークで推論ワークロードを高速化する場合に適しています。Dataflow は、各ワーカーが 1 つ以上の TPU デバイスを管理するシングルホスト TPU 構成をサポートしています。

詳細については、Dataflow での TPU のサポートをご覧ください。

ワークフロー オーケストレーション

以降のセクションでは、ワークフロー オーケストレーションのユースケースについて説明します。

TFXKubeflow Pipelines(KFP)はどちらも Apache Beam コンポーネントを使用します。

Vertex AI Pipelines で Dataflow を使用する場合

Vertex AI Pipelines では、ML ワークフローをサーバーレスな方法でオーケストレートすることで、ML システムの自動化、モニタリング、管理を行うことができます。Vertex AI Pipelines では、TFX または KFP で定義されたワークフロー有向非巡回グラフ(DAG)をオーケストレートし、Vertex ML Metadata を使用して ML アーティファクトを自動的に追跡できます。Dataflow を TFX や KFP に組み込む方法については、次のセクションの情報をご覧ください。

Kubeflow Pipelines で Dataflow を使用する場合

Kubeflow は、Kubernetes での ML ワークフローのデプロイを使いやすく移植可能、かつスケーラブルなものすることに特化した ML ツールキットです。Kubeflow Pipelines は、Kubeflow Pipelines SDK を使用して構築された再利用可能なエンドツーエンドの ML ワークフローです。

Kubeflow Pipelines サービスは、エンドツーエンドのオーケストレーションを提供し、テストと再利用を簡単にすることを目的としています。KFP では、オーケストレーション手法のテストやテストの管理を行うことができます。毎回最初からやり直すことなく、コンポーネントとパイプラインを再利用して、複数のエンドツーエンドのソリューションを作成できます。

KFP で Dataflow を使用する場合は、DataflowPythonJobOP 演算子または DataflowFlexTemplateJobOp 演算子を使用できます。完全なカスタム コンポーネントを作成することもできます。DataflowPythonJobOP 演算子を使用することをおすすめします。

完全なカスタム コンポーネントを作成する場合は、Vertex AI ドキュメントの Dataflow コンポーネントのページをご覧ください。

TFX で Dataflow を使用する場合

TFX パイプライン コンポーネントは TFX ライブラリ上に構築されます。データ処理ライブラリは Apache Beam を直接使用します。たとえば、TensorFlow Transform はユーザーの呼び出しを Apache Beam に変換します。したがって、追加の構成作業を行うことなく、TFX パイプラインで Apache Beam と Dataflow を使用できます。Dataflow で TFX を使用するには、TFX パイプラインを構築するときに Dataflow ランナーを使用します。詳しくは、次のリソースをご覧ください。