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
,
output_dim
)
def
forward
(
self
,
x
):
out
=
self
.
linear
(
x
)
return
out
Prepare train and test data for an example model
This example model is a 5 times table.
-
xcontains values in the range from 0 to 99. -
yis a list of 5 *x. -
value_to_predictincludes values outside of the training data.
x
=
numpy
.
arange
(
0
,
100
,
dtype
=
numpy
.
float32
)
.
reshape
(
-
1
,
1
)
y
=
(
x
*
5
)
.
reshape
(
-
1
,
1
)
value_to_predict
=
numpy
.
array
([
20
,
40
,
60
,
90
],
dtype
=
numpy
.
float32
)
.
reshape
(
-
1
,
1
)
Train the linear regression mode on 5 times data
Use the following code to train your linear regression model on the 5 times table.
five_times_model
=
LinearRegression
()
optimizer
=
torch
.
optim
.
Adam
(
five_times_model
.
parameters
())
loss_fn
=
torch
.
nn
.
L1Loss
()
"""
Train the five_times_model
"""
epochs
=
10000
tensor_x
=
torch
.
from_numpy
(
x
)
tensor_y
=
torch
.
from_numpy
(
y
)
for
epoch
in
range
(
epochs
):
y_pred
=
five_times_model
(
tensor_x
)
loss
=
loss_fn
(
y_pred
,
tensor_y
)
five_times_model
.
zero_grad
()
loss
.
backward
()
optimizer
.
step
()
Save the model using torch.save()
, and then confirm that the saved model file exists.
torch
.
save
(
five_times_model
.
state_dict
(),
save_model_dir_multiply_five
)
print
(
os
.
path
.
exists
(
save_model_dir_multiply_five
))
# Verify that the model is saved.
True
Prepare train and test data for a 10 times model
This example model is a 10 times table.
-
xcontains values in the range from 0 to 99. -
yis a list of 10 *x.
x
=
numpy
.
arange
(
0
,
100
,
dtype
=
numpy
.
float32
)
.
reshape
(
-
1
,
1
)
y
=
(
x
*
10
)
.
reshape
(
-
1
,
1
)
Train the linear regression model on 10 times data
Use the following to train your linear regression model on the 10 times table.
ten_times_model
=
LinearRegression
()
optimizer
=
torch
.
optim
.
Adam
(
ten_times_model
.
parameters
())
loss_fn
=
torch
.
nn
.
L1Loss
()
epochs
=
10000
tensor_x
=
torch
.
from_numpy
(
x
)
tensor_y
=
torch
.
from_numpy
(
y
)
for
epoch
in
range
(
epochs
):
y_pred
=
ten_times_model
(
tensor_x
)
loss
=
loss_fn
(
y_pred
,
tensor_y
)
ten_times_model
.
zero_grad
()
loss
.
backward
()
optimizer
.
step
()
Save the model using torch.save()
.
torch
.
save
(
ten_times_model
.
state_dict
(),
save_model_dir_multiply_ten
)
print
(
os
.
path
.
exists
(
save_model_dir_multiply_ten
))
# verify if the model is saved
True
Pattern 1: RunInference for predictions
This pattern demonstrates how to use RunInference for predictions.
Use RunInference within the pipeline
- Create a PyTorch model handler object by passing required arguments such as
state_dict_path,model_class, andmodel_paramsto thePytorchModelHandlerTensorclass. - Pass the
PytorchModelHandlerTensorobject to the RunInference transform to perform predictions on unkeyed data.
torch_five_times_model_handler
=
PytorchModelHandlerTensor
(
state_dict_path
=
save_model_dir_multiply_five
,
model_class
=
LinearRegression
,
model_params
=
{
'input_dim'
:
1
,
'output_dim'
:
1
}
)
pipeline
=
beam
.
Pipeline
()
with
pipeline
as
p
:
(
p
|
"ReadInputData"
>> beam
.
Create
(
value_to_predict
)
|
"ConvertNumpyToTensor"
>> beam
.
Map
(
torch
.
Tensor
)
|
"RunInferenceTorch"
>> RunInference
(
torch_five_times_model_handler
)
|
beam
.
Map
(
print
)
)
PredictionResult(example=tensor([20.]), inference=tensor([102.0095], grad_fn=<UnbindBackward0>)) PredictionResult(example=tensor([40.]), inference=tensor([201.2056], grad_fn=<UnbindBackward0>)) PredictionResult(example=tensor([60.]), inference=tensor([300.4017], grad_fn=<UnbindBackward0>)) PredictionResult(example=tensor([90.]), inference=tensor([449.1959], grad_fn=<UnbindBackward0>))
Pattern 2: Postprocess RunInference results
This pattern demonstrates how to postprocess the RunInference results.
Add a PredictionProcessor
to the pipeline after RunInference
. PredictionProcessor
processes the output of the RunInference
transform.
class
PredictionProcessor
(
beam
.
DoFn
):
"""
A processor to format the output of the RunInference transform.
"""
def
process
(
self
,
element
:
PredictionResult
):
input_value
=
element
.
example
output_value
=
element
.
inference
yield
(
f
"input is
{
input_value
.
item
()
}
output is
{
output_value
.
item
()
}
"
)
pipeline
=
beam
.
Pipeline
()
with
pipeline
as
p
:
(
p
|
"ReadInputData"
>> beam
.
Create
(
value_to_predict
)
|
"ConvertNumpyToTensor"
>> beam
.
Map
(
torch
.
Tensor
)
|
"RunInferenceTorch"
>> RunInference
(
torch_five_times_model_handler
)
|
"PostProcessPredictions"
>> beam
.
ParDo
(
PredictionProcessor
())
|
beam
.
Map
(
print
)
)
input is 20.0 output is 102.00947570800781 input is 40.0 output is 201.20559692382812 input is 60.0 output is 300.4017028808594 input is 90.0 output is 449.1958923339844
Pattern 3: Attach a key
This pattern demonstrates how attach a key to allow your model to handle keyed data.
Modify the model handler and post processor
Modify the pipeline to read from sources like CSV files and BigQuery.
In this step, you take the following actions:
- To handle keyed data, wrap the
PytorchModelHandlerTensorobject aroundKeyedModelHandler. - Add a map transform that converts a table row into
Tuple[str, float]. - Add a map transform that converts
Tuple[str, float]toTuple[str, torch.Tensor]. - Modify the post-inference processor to output results with the key.
class
PredictionWithKeyProcessor
(
beam
.
DoFn
):
def
__init__
(
self
):
beam
.
DoFn
.
__init__
(
self
)
def
process
(
self
,
element
:
Tuple
[
str
,
PredictionResult
]):
key
=
element
[
0
]
input_value
=
element
[
1
]
.
example
output_value
=
element
[
1
]
.
inference
yield
(
f
"key:
{
key
}
, input:
{
input_value
.
item
()
}
output:
{
output_value
.
item
()
}
"
)
Create a source with attached key
This section shows how to create either a BigQuery or a CSV source with an attached key.
Use BigQuery as the source
Follow these steps to use BigQuery as your source.
To install the Google Cloud BigQuery API, use pip
:
%
pip
install
--
upgrade
google
-
cloud
-
bigquery
--
quiet
Create a table in BigQuery using the following snippet, which has two columns. The first column holds the key and the second column holds the test value. To use BiqQuery, you need a Google Cloud account with the BigQuery API enabled.
gcloud
config
set
project
$project
Updated property [core/project].
from
google.cloud
import
bigquery
client
=
bigquery
.
Client
(
project
=
project
)
# Make sure the dataset_id is unique in your project.
dataset_id
=
'
{project}
.maths'
.
format
(
project
=
project
)
dataset
=
bigquery
.
Dataset
(
dataset_id
)
# Modify the location based on your project configuration.
dataset
.
location
=
'US'
dataset
=
client
.
create_dataset
(
dataset
,
exists_ok
=
True
)
# Table name in the BigQuery dataset.
table_name
=
'maths_problems_1'
query
=
"""
CREATE OR REPLACE TABLE
{project}
.maths.
{table}
( key STRING OPTIONS(description="A unique key for the maths problem"),
value FLOAT64 OPTIONS(description="Our maths problem" ) );
INSERT INTO maths.
{table}
VALUES
("first_question", 105.00),
("second_question", 108.00),
("third_question", 1000.00),
("fourth_question", 1013.00)
"""
.
format
(
project
=
project
,
table
=
table_name
)
create_job
=
client
.
query
(
query
)
create_job
.
result
()
<google.cloud.bigquery.table._EmptyRowIterator at 0x7f5694206650>
To read keyed data, use BigQuery as the pipeline source.
pipeline_options
=
PipelineOptions
()
.
from_dictionary
({
'temp_location'
:
f
'gs://
{
bucket
}
/tmp'
,
})
pipeline
=
beam
.
Pipeline
(
options
=
pipeline_options
)
keyed_torch_five_times_model_handler
=
KeyedModelHandler
(
torch_five_times_model_handler
)
table_name
=
'maths_problems_1'
table_spec
=
f
'
{
project
}
:maths.
{
table_name
}
'
with
pipeline
as
p
:
(
p
|
"ReadFromBQ"
>> beam
.
io
.
ReadFromBigQuery
(
table
=
table_spec
)
|
"PreprocessData"
>> beam
.
Map
(
lambda
x
:
(
x
[
'key'
],
x
[
'value'
]))
|
"ConvertNumpyToTensor"
>> beam
.
Map
(
lambda
x
:
(
x
[
0
],
torch
.
Tensor
([
x
[
1
]])))
|
"RunInferenceTorch"
>> RunInference
(
keyed_torch_five_times_model_handler
)
|
"PostProcessPredictions"
>> beam
.
ParDo
(
PredictionWithKeyProcessor
())
|
beam
.
Map
(
print
)
)
key: third_question, input: 1000.0 output: 4962.61962890625 key: second_question, input: 108.0 output: 538.472412109375 key: first_question, input: 105.0 output: 523.5929565429688 key: fourth_question, input: 1013.0 output: 5027.0966796875
Use a CSV file as the source
Follow these steps to use a CSV file as your source.
Create a CSV file with two columns: one named key
that holds the keys, and a second named value
that holds the test values.
# creates a CSV file with the values.
csv_values
=
[(
"first_question"
,
105.00
),
(
"second_question"
,
108.00
),
(
"third_question"
,
1000.00
),
(
"fourth_question"
,
1013.00
)]
input_csv_file
=
"./maths_problem.csv"
with
open
(
input_csv_file
,
'w'
)
as
f
:
writer
=
csv
.
writer
(
f
)
writer
.
writerow
([
'key'
,
'value'
])
for
row
in
csv_values
:
writer
.
writerow
(
row
)
assert
os
.
path
.
exists
(
input_csv_file
)
==
True
pipeline_options
=
PipelineOptions
()
.
from_dictionary
({
'temp_location'
:
f
'gs://
{
bucket
}
/tmp'
,
})
pipeline
=
beam
.
Pipeline
(
options
=
pipeline_options
)
keyed_torch_five_times_model_handler
=
KeyedModelHandler
(
torch_five_times_model_handler
)
with
pipeline
as
p
:
df
=
p
|
beam
.
dataframe
.
io
.
read_csv
(
input_csv_file
)
pc
=
to_pcollection
(
df
)
(
pc
|
"ConvertNumpyToTensor"
>> beam
.
Map
(
lambda
x
:
(
x
[
0
],
torch
.
Tensor
([
x
[
1
]])))
|
"RunInferenceTorch"
>> RunInference
(
keyed_torch_five_times_model_handler
)
|
"PostProcessPredictions"
>> beam
.
ParDo
(
PredictionWithKeyProcessor
())
|
beam
.
Map
(
print
)
)
key: first_question, input: 105.0 output: 523.5929565429688 key: second_question, input: 108.0 output: 538.472412109375 key: third_question, input: 1000.0 output: 4962.61962890625 key: fourth_question, input: 1013.0 output: 5027.0966796875
Pattern 4: Inference with multiple models in the same pipeline
This pattern demonstrates how use inference with multiple models in the same pipeline.
Multiple models in parallel
This section demonstrates how use inference with multiple models in parallel.
Create a torch model handler for the 10 times model using PytorchModelHandlerTensor
.
torch_ten_times_model_handler
=
PytorchModelHandlerTensor
(
state_dict_path
=
save_model_dir_multiply_ten
,
model_class
=
LinearRegression
,
model_params
=
{
'input_dim'
:
1
,
'output_dim'
:
1
}
)
keyed_torch_ten_times_model_handler
=
KeyedModelHandler
(
torch_ten_times_model_handler
)
In this section, the same data is run through two different models: the one that we use to multiply by 5 and a new model that learns to multiply by 10.
pipeline_options
=
PipelineOptions
()
.
from_dictionary
(
{
'temp_location'
:
f
'gs://
{
bucket
}
/tmp'
})
pipeline
=
beam
.
Pipeline
(
options
=
pipeline_options
)
read_from_bq
=
beam
.
io
.
ReadFromBigQuery
(
table
=
table_spec
)
with
pipeline
as
p
:
multiply_five
=
(
p
|
read_from_bq
|
"CreateMultiplyFiveTuple"
>> beam
.
Map
(
lambda
x
:
(
'
{}
{}
'
.
format
(
x
[
'key'
],
'* 5'
),
x
[
'value'
]))
|
"ConvertNumpyToTensorFiveTuple"
>> beam
.
Map
(
lambda
x
:
(
x
[
0
],
torch
.
Tensor
([
x
[
1
]])))
|
"RunInferenceTorchFiveTuple"
>> RunInference
(
keyed_torch_five_times_model_handler
)
)
multiply_ten
=
(
p
|
read_from_bq
|
"CreateMultiplyTenTuple"
>> beam
.
Map
(
lambda
x
:
(
'
{}
{}
'
.
format
(
x
[
'key'
],
'* 10'
),
x
[
'value'
]))
|
"ConvertNumpyToTensorTenTuple"
>> beam
.
Map
(
lambda
x
:
(
x
[
0
],
torch
.
Tensor
([
x
[
1
]])))
|
"RunInferenceTorchTenTuple"
>> RunInference
(
keyed_torch_ten_times_model_handler
)
)
inference_result
=
((
multiply_five
,
multiply_ten
)
|
beam
.
Flatten
()
|
beam
.
ParDo
(
PredictionWithKeyProcessor
()))
inference_result
|
beam
.
Map
(
print
)
key: third_question * 10, input: 1000.0 output: 9889.59765625 key: second_question * 10, input: 108.0 output: 1075.4891357421875 key: first_question * 10, input: 105.0 output: 1045.84521484375 key: fourth_question * 10, input: 1013.0 output: 10018.0546875 key: third_question * 5, input: 1000.0 output: 4962.61962890625 key: second_question * 5, input: 108.0 output: 538.472412109375 key: first_question * 5, input: 105.0 output: 523.5929565429688 key: fourth_question * 5, input: 1013.0 output: 5027.0966796875
Multiple models in sequence
This section demonstrates how use inference with multiple models in sequence.
In a sequential pattern, data is sent to one or more models in sequence, with the output from each model chaining to the next model. The following list demonstrates the sequence used in this section.
- Read the data from BigQuery.
- Map the data.
- Use RunInference with the multiply by 5 model.
- Process the results.
- Use RunInference with the multiply by 10 model.
- Process the results.
def
process_interim_inference
(
element
):
key
,
prediction_result
=
element
input_value
=
prediction_result
.
example
inference
=
prediction_result
.
inference
formatted_input_value
=
'original input is `
{}
{}
`'
.
format
(
key
,
input_value
)
return
formatted_input_value
,
inference
pipeline_options
=
PipelineOptions
()
.
from_dictionary
(
{
'temp_location'
:
f
'gs://
{
bucket
}
/tmp'
})
pipeline
=
beam
.
Pipeline
(
options
=
pipeline_options
)
with
pipeline
as
p
:
multiply_five
=
(
p
|
beam
.
io
.
ReadFromBigQuery
(
table
=
table_spec
)
|
"CreateMultiplyFiveTuple"
>> beam
.
Map
(
lambda
x
:
(
x
[
'key'
],
x
[
'value'
]))
|
"ConvertNumpyToTensorFiveTuple"
>> beam
.
Map
(
lambda
x
:
(
x
[
0
],
torch
.
Tensor
([
x
[
1
]])))
|
"RunInferenceTorchFiveTuple"
>> RunInference
(
keyed_torch_five_times_model_handler
)
)
inference_result
=
(
multiply_five
|
"ExtractResult"
>> beam
.
Map
(
process_interim_inference
)
|
"RunInferenceTorchTenTuple"
>> RunInference
(
keyed_torch_ten_times_model_handler
)
|
beam
.
ParDo
(
PredictionWithKeyProcessor
())
)
inference_result
|
beam
.
Map
(
print
)
key: original input is `third_question tensor([1000.])`, input: 4962.61962890625 output: 49045.37890625 key: original input is `second_question tensor([108.])`, input: 538.472412109375 output: 5329.11083984375 key: original input is `first_question tensor([105.])`, input: 523.5929565429688 output: 5182.08251953125 key: original input is `fourth_question tensor([1013.])`, input: 5027.0966796875 output: 49682.49609375
