AI & 機械学習

Apache Beam の新しい Dataflow ML イノベーションで ML モデルの実行がより簡単に

2023年7月1日

Google Cloud Japan Team

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

Harvard Business Review によると、AI 用の自社モデルを本番環境に投入できている企業はわずか 20% です。Google Cloud Dataflow は、最もよく利用されているオープンソースフレームワークの一つである Apache Beam を基盤として活用します。Apache Beam は、統合型プログラミングモデルであり、バッチおよびストリーミングパイプラインの開発向け SDK を提供します。オープンなプロダクトを構築するための取り組みを続けるなかで、Google は Beam コミュニティと密に連携して、ML に焦点を当てた 3 つの新機能を Dataflow に追加することになりました。これらの機能とコミュニティの膨大なリソースを併せて活用すれば、本番環境における大規模なストリーミング ML モデルの実行を簡略化できます。

自動モデル更新: データや ML モデルは常に変化しているため、効果を維持するには、モデルを継続的に再トレーニングして改善する必要があります。こうしたモデルの更新では、データエンジニアリングチームが再デプロイに膨大な労力を注ぐ必要はありません。新しいストリーミングの自動モデル更新機能を使用すると、モデルを更新し、データストリーム処理を一時停止することなく稼働中のストリーミングパイプラインでホットスワップできるため、ダウンタイムの回避が可能となります。
TensorFlow Hub のインテグレーション: TensorFlow Hub は、事前トレーニング済み ML モデルのリポジトリであり、素早く簡単にアプリケーションに ML モデルをデプロイしたい研究者やデベロッパーの貴重なリソースとなっています。オープンソースモデルは、単体でも複数の組み合わせ（アンサンブル）でも、データサイエンティストによく使用されるツールです。たとえば、お客様が投稿したレビューの感情分析を行う際、カスタムのドメイン固有モデルにデータを渡す前に、エンベディングの段階でオープンソースのモデルを使用します。Apache Beam ユーザーがこの手順を簡単に実施できるように、TensorFlow Hub のインテグレーションによって、わずか数行のコード入力でモデルをダウンロードしてバッチまたはストリーミング Dataflow パイプライン内で利用できるようにしました。
モデルアンサンブル: ML のフレームワークが急増する中、使用できる処理パイプラインが 1 つに制限されてしまうのは理想的とはいえません。Apache Beam SDK を活用すると、TensorFlow、PyTorch、SKLearn などの複数のフレームワークを、アンサンブルとして 1 つのパイプラインで使用することが可能になります。Apache Beam コミュニティは最近、ONNX モデルハンドラのサポートを開始し、サポート対象のフレームワークがさらに増えました。

ストリーミングの自動モデル更新

Apache Beam は、単一の統合型 API を使用してストリーミングとバッチの両方に対応できることで知られています。この機能によりデベロッパーは、コードの変更や多数の新しい import ステートメントの追加を必要とせずに、バッチとストリーミングを切り替えることができます。しかしこれまでは、モデルが再トレーニングされるたびに、オペレーションエンジニアはパイプラインのさまざまなライフサイクルイベントを処理する必要がありました。この自動モデル更新機能では、パイプライン全体の更新や再デプロイを必要とせず、RunInference 変換が自動的にモデルハンドラを更新します。

この機能では、次の 2 つのモードが利用可能です。

ウォッチモード: このモードでは、Google Cloud Storage から更新を pull します。RunInference の WatchFilePattern クラスを使用して、最新のファイルを確認します。その後、タイムスタンプに基づいて file_pattern を照合することで更新します。最新の ModelMetadata が出力され、これを RunInference の PTransform で使用して、Beam パイプラインを止めることなく ML モデルを更新します。
イベントモード: Pub/Sub のような無限ソースにパイプラインを接続することで、更新イベントが変換に直接送られ、オンデマンドでモデルのホットスワップが開始されます。

TensorFlow Hub のインテグレーション

新しいネイティブの TensorFlow モデルハンドラでは、引数として URL をモデルに渡すことで、Dataflow パイプラインで TensorFlow Hub モデルを使用できます。TensorFlow Hub は、TensorFlow フレームワークや、その高水準なラッパーである Keras 内で使用可能な、事前トレーニング済み ML モデルのリポジトリです。パワフルなモデルをトレーニングなしで、さまざまなタスクに活用できます。リポジトリには 1,000 を超えるモデルが含まれており、コミュニティの協力によりさらに増え続けています。

それでは、小売店舗管理や e コマースサイトなど、さまざまな業界でよく使用されている画像分類の例を見てみましょう。次の例では、tfhub 上で mobilenet_v2 モデルを使用しています。わずか数行のコードで、数百万の画像処理が可能なモデルをパイプラインで使用できます（このノートブックの実行可能なコピーは、こちらからご利用いただけます）。

読み込んでいます...

from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerTensor
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RunInference
from typing import Iterable

model_handler = TFModelHandlerTensor(model_uri="https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4")

class PostProcessor(beam.DoFn):
  """PredictionResult を処理して予測ラベルを取得。
 予測ラベルを返す。
  """
  def setup(self):
    labels_path = tf.keras.utils.get_file(
        'ImageNetLabels.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'
    )
    self._imagenet_labels = np.array(open(labels_path).read().splitlines())

def process(self, element: PredictionResult) -> Iterable[str]:
    predicted_class = np.argmax(element.inference)
    predicted_class_name = self._imagenet_labels[predicted_class]
    yield "Predicted Label: {}".format(predicted_class_name.title())

with beam.Pipeline() as p:
  _ = (p
    | "Create PCollection" >> beam.Create([img_tensor])
    | "Perform inference" >> RunInference(model_handler)
    | "Post Processing" >> beam.ParDo(PostProcessor())
    | "Print" >> beam.Map(print))

このコードを数百万の画像用に機能させるためには、ソースをオブジェクトストア上のバケットに変更し、多数の Apache Beam シンクの一つを使用して目的の場所に結果を出力します。

モデルアンサンブル

1 つのビジネスユースケースであったとしても、複数のフレームワークを使用して構築するケースがますます増えています。そのため、可能な限り多くのフレームワークを RunInference でもサポートすることが必要になってきます。RunInference では、モデルハンドラにこうしたサポートを追加しています。

今回のリリース前までは、Apache Beam では以下をサポートしていました。

Beam コミュニティは、ONNX のモデルハンドラを伴う ONNX runtime を追加しました。パイプラインのサンプルは、onnx_sentiment_classification.py で確認できます。

1 つのパイプラインでの複数モデルの使用について詳しくは、こちらをご覧ください。BLIP モデルと CLIP モデルの併用について説明するノートブックも含まれています。