Make predictions with PyTorch models in ONNX format

Overview

Open Neural Network Exchange (ONNX) provides a uniform format designed to represent any machine learning frameworks. BigQuery ML support for ONNX allows you to:

• Train a model using your favorite framework.
• Convert the model into ONNX model format. For more information, see Converting to ONNX format.
• Import the ONNX model into BigQuery and make predictions using BigQuery ML.

This tutorial shows you how to import ONNX models trained with PyTorch into a BigQuery dataset and use them to make predictions from a SQL query. You can import ONNX models using these interfaces:

• The Google Cloud console
• The bq query command in the bq command-line tool
• The BigQuery API

For more information about importing ONNX models into BigQuery, including format and storage requirements, see The CREATE MODEL statement for importing ONNX models.

Objectives

In this tutorial, you will:

• Create and train models with PyTorch.
• Convert the models to ONNX format using torch.onnx.
• Import the ONNX models into BigQuery and make predictions.

Create a PyTorch vision model for image classification

Import a PyTorch pretrained resnet18 that accepts decoded image data returned by the BigQuery ML ML.DECODE_IMAGE and ML.RESIZE_IMAGE functions.

import torch
import torch.nn as nn

# Define model input format to match the output format of
# ML.DECODE_IMAGE function: [height, width, channels]
dummy_input = torch.randn(1, 224, 224, 3, device="cpu")

# Load a pretrained pytorch model for image classification
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)

# Reshape input format from [batch_size, height, width, channels]
# to [batch_size, channels, height, width]
class ReshapeLayer(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x.permute(0, 3, 1, 2)  # reorder dimensions
        return x

class ArgMaxLayer(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
       return torch.argmax(x, dim=1)

final_model = nn.Sequential(
    ReshapeLayer(),
    model,
    nn.Softmax(),
    ArgMaxLayer()
)

Convert the model into ONNX format and save

Use torch.onnx to export the PyTorch vision model to an ONNX file named resnet18.onnx.

torch.onnx.export(final_model,            # model being run
                  dummy_input,            # model input
                  "resnet18.onnx",        # where to save the model
                  opset_version=10,       # the ONNX version to export the model to
                  input_names = ['input'],         # the model's input names
                  output_names = ['class_label'])  # the model's output names

Upload the ONNX model to Cloud Storage

Create a Cloud Storage bucket to store the ONNX model file, and then upload the saved ONNX model file to your Cloud Storage bucket. For more information, see Upload objects from a filesystem.

Import the ONNX model into BigQuery

This step assumes you have uploaded the ONNX model to your Cloud Storage bucket. An example model is stored at gs://cloud-samples-data/bigquery/ml/onnx/resnet18.onnx.

bq query \
--use_legacy_sql=false \
"CREATE MODEL
  `mydataset.mymodel`
OPTIONS
  (MODEL_TYPE='ONNX',
   MODEL_PATH='gs://bucket/path/to/onnx_model/*')"

bq query --use_legacy_sql=false \
"CREATE OR REPLACE MODEL
  `example_dataset.imported_onnx_model`
OPTIONS
  (MODEL_TYPE='ONNX',
   MODEL_PATH='gs://cloud-samples-data/bigquery/ml/onnx/resnet18.onnx')"

$ bq ls example_dataset

       tableId          Type    Labels   Time Partitioning
 --------------------- ------- -------- -------------------
  imported_onnx_model   MODEL

{
  "query": "CREATE MODEL `project_id:mydataset.mymodel` OPTIONS(MODEL_TYPE='ONNX' MODEL_PATH='gs://bucket/path/to/onnx_model/*')"
}

Create an object table in BigQuery to access image data

To access unstructured data in BigQuery, you need to create an object table. See Create object tables for detailed instructions.

Create an object table named goldfish_image_table on a goldfish image stored at gs://mybucket/goldfish.jpg.

CREATE EXTERNAL TABLE `example_dataset.goldfish_image_table`
WITH CONNECTION `us.my-connection`
OPTIONS(
  object_metadata = 'SIMPLE',
  uris = ['gs://mybucket/goldfish.jpg'],
  max_staleness = INTERVAL 1 DAY,
  metadata_cache_mode = 'AUTOMATIC');

Make predictions with the imported ONNX model

bq query \
--use_legacy_sql=false \
'SELECT *
 FROM ML.PREDICT(
   MODEL `my_project.my_dataset.my_model`,
   (SELECT * FROM input_data))'

{
  "query": "SELECT * FROM ML.PREDICT(MODEL `my_project.my_dataset.my_model`, (SELECT * FROM input_data))"
}

What's next

• For more information about importing ONNX models, see The CREATE MODEL statement for ONNX models.
• For more information about available ONNX converters and tutorials, see Converting to ONNX format.
• For an overview of BigQuery ML, see Introduction to BigQuery ML.
• To get started using BigQuery ML, see Create machine learning models in BigQuery ML.
• For more information about working with models, see these resources:
  • Get model metadata
  • Update model metadata
  • Manage models
  • Delete models