Apache Beam RunInference for PyTorch

Run in Google Colab View source on GitHub

This notebook demonstrates the use of the RunInference transform for PyTorch. Apache Beam includes implementations of the ModelHandler class for users of PyTorch. For more information about using RunInference, see Get started with AI/ML pipelines in the Apache Beam documentation.

This notebook illustrates common RunInference patterns, such as:

  • Using a database with RunInference.
  • Postprocessing results after using RunInference.
  • Inference with multiple models in the same pipeline.

The linear regression models used in these samples are trained on data that correspondes to the 5 and 10 times tables; that is, y = 5x and y = 10x respectively.

Dependencies

The RunInference library is available in Apache Beam versions 2.40 and later.

To use Pytorch RunInference API, you need to install the PyTorch module. To install PyTorch, use pip:

pip install apache_beam[interactive,gcp,dataframe] --quiet
 
%pip install torch --quiet

import argparse
import csv
import json
import os
import torch
from typing import Tuple

import apache_beam as beam
import numpy
from apache_beam.io.gcp.bigquery import ReadFromBigQuery
from apache_beam.ml.inference.base import KeyedModelHandler
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RunInference
from apache_beam.dataframe.convert import to_pcollection
from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor
from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor
from apache_beam.options.pipeline_options import PipelineOptions

import warnings
warnings.filterwarnings('ignore')

from google.colab import auth
auth.authenticate_user()

# Constants
project = "<your GCP project>" # @param {type:'string'}
bucket = "<your GCP bucket>" # @param {type:'string'}

# To avoid warnings, set the project.
os.environ['GOOGLE_CLOUD_PROJECT'] = project

save_model_dir_multiply_five = 'five_times_table_torch.pt'
save_model_dir_multiply_ten = 'ten_times_table_torch.pt'

Create data and PyTorch models for the RunInference transform

Create linear regression models, prepare train and test data, and train models.

Create a linear regression model in PyTorch

Use the following code to create a linear regression model.

class LinearRegression(torch.nn.Module):
    def __init__(self, input_dim=1, output_dim=1):
        super().__init__()
        self.linear = torch.nn.Linear(input_dim,