Dataflow の Java パイプラインからの推論ワークロードの実行方法

※この投稿は米国時間 2023 年 8 月 1 日に、Google Cloud blog に投稿されたものの抄訳です。

Python ベースの ML ワークロードをより簡単に Java または Go パイプラインから実行できる機能が多くのデータ エンジニアから求められています。Java または Go で記述された既存のパイプラインがあり、それを ML で強化したい、Java または Go で記述された特別なソースや変換があるのでそれを利用したい、あるいは単純に Java または Go での作業が気に入っているなど、その理由はさまざまですが、理由はどうあれ、Dataflow の多言語機能を使用すると、Python の ML ワークロードを簡単に実行することができるようになりました。

Dataflow で多言語による推論を実行するには、主に 2 つの方法があります。

1）Java または Go の RunInference 変換を使用して実行する。TensorFlow、PyTorch、Sklearn などのフレームワークを使用して開発された Python モデルが存在し、そのモデルを使用して推論を実行したいが、前処理や後処理の大部分では Java または Go を使用したい場合、Java または Go の RunInference 変換を使用できます。Dataflow は、バックグラウンドで自動的に Python ワーカーをスピンアップし、推論を実行します。

入力のバッチ処理、効率的なモデルの共有、リモート ファイル システムまたはモデルハブからのモデルの読み込みなどのボイラープレート コードの大部分は、RunInference 変換により自動的に処理されます。Beam 2.47 の時点で、RunInference では TensorFlow、PyTorch、Sklearn、XGBoost、ONNX、および TensorRT モデルが完全にサポートされています。独自のカスタム モデルハンドラを構築して他のフレームワークをサポートすることもできます。Java RunInference 変換を使用したパイプライン全体の例については、https://github.com/apache/beam/blob/master/examples/java/src/main/java/org/apache/beam/examples/multilanguage/SklearnMnistClassification.java をご覧ください。

2）前処理と後処理、および推論のすべてを実行するカスタム Python 変換を作成する（この場合でも RunInference 変換を利用できます）。これは外部変換（PythonExternalTransform）を使用して呼び出すことができます。この方法を使用すると、前処理および後処理でお好みの Python ライブラリを使用できます。

外部変換を使用したサンプルのチュートリアル全体については、https://beam.apache.org/documentation/ml/multi-language-inference/ をご覧ください。

Dataflow を使用すれば、Java または Go から大規模な推論を簡単に実行できます。これを利用したパイプラインにより、既存のロジックを根本から見直すことなく、Java API を引き続き使用しながら、Python ML フレームワークの潜在能力をいかんなく発揮させることができます。

詳細

詳しく学習したい場合は、apache/beam リポジトリに用意されている Java パイプラインから推論を実行するエンドツーエンドのサンプルを通して理解を深めることができます。

- Dataflow ML シニア ソフトウェア エンジニア Danny McCormick