Dataflow ML

Dataflow ML allows you to use Dataflow to deploy and manage end-to-end machine learning (ML) pipelines. Use ML models to do local and remote inference with batch and streaming pipelines. Use data processing tools to prepare your data for model training and to process the results of the models.

About Dataflow ML

with pipeline as p:
  predictions = (
      p
      | beam.ReadFromSource('a_source')
      | RunInference(MODEL_HANDLER))

RunInference transform

Using RunInference is as straightforward as adding the transform code to your pipeline. In this example, MODEL_HANDLER is the model configuration object.

Prediction and inference

Whether you want to classify images in real-time, run remote inference calls, or build a custom model handler, you can find end-to-end Dataflow ML examples.

See use cases

Data processing

Data processing for ML training is often used in the context of ML-OPS frameworks. These frameworks orchestrate ML workflows, which contain various pre- and post-processing steps. Read more about how Dataflow integrates with these systems.

Learn more

Prediction and inference use cases

I have a Pytorch model

Use a pre-trained model with Pytorch.

I have a scikit-learn model

Use a pre-trained model with scikit-learn.

I have a TensorFlow model

Use a pre-trained model with Tensorflow.

I have a model deployed on Vertex AI

Vertex AI supports both online and batch prediction and has built-in integration with other ML workflows, like model registry, pipelines, model monitoring, and Explainable AI. Run inference using a model deployed on Vertex AI by using custom inference calls.

See more

Explore all the available prediction and inference use cases.

Dataflow ML in ML workflows

I want to use Dataflow with Vertex AI Pipelines

Vertex AI Pipelines helps you to automate, monitor, and govern your ML systems by orchestrating your ML workflows in a serverless manner. Use Vertex AI Pipelines to orchestrate workflow DAGs defined by either TFX or KFP and to automatically track your ML artifacts using Vertex ML Metadata.

I want to use Dataflow with TFX

TensorFlow Extended (TFX) lets you deploy end-to-end ML pipelines by using an orchestration framework that has a built-in integration with Apache Beam and the Dataflow runner.

I want to use Dataflow with KFP

Kubeflow makes deployments of ML workflows on Kubernetes simple, portable, and scalable. Kubeflow Pipelines are reusable end-to-end ML workflows built using the Kubeflow Pipelines SDK.