Make predictions with scikit-learn models in ONNX format

This tutorial shows you how to import an Open Neural Network Exchange (ONNX) model that's trained with scikit-learn. You import the model into a BigQuery dataset and use it to make predictions using a SQL query.

ONNX provides a uniform format that is designed to represent any machine learning (ML) framework. BigQuery ML support for ONNX lets you do the following:

• Train a model using your favorite framework.
• Convert the model into the ONNX model format.
• Import the ONNX model into BigQuery and make predictions using BigQuery ML.

Objectives

• Create and train a model using scikit-learn.
• Convert the model to ONNX format using sklearn-onnx.
• Use the CREATE MODEL statement to import the ONNX model into BigQuery.
• Use the ML.PREDICT function to make predictions with the imported ONNX model.

Costs

In this document, you use the following billable components of Google Cloud:

• BigQuery
• BigQuery ML
• Cloud Storage

To generate a cost estimate based on your projected usage, use the pricing calculator.

When you finish the tasks that are described in this document, you can avoid continued billing by deleting the resources that you created. For more information, see Clean up.

Console

1. In the Google Cloud console, go to the BigQuery page.

   Go to the BigQuery page

2. In the Explorer pane, click your project name.

3. Click more_vert View actions > Create dataset

4. On the Create dataset page, do the following:

   • For Dataset ID, enter bqml_tutorial.

   • For Location type, select Multi-region, and then select US.

   • Leave the remaining default settings as they are, and click Create dataset.

bq

To create a new dataset, use the bq mk --dataset command.

1. Create a dataset named bqml_tutorial with the data location set to US.

   bq mk --dataset \
     --location=US \
     --description "BigQuery ML tutorial dataset." \
     bqml_tutorial

2. Confirm that the dataset was created:

   bq ls

API

Call the datasets.insert method with a defined dataset resource.

{
  "datasetReference": {
     "datasetId": "bqml_tutorial"
  }
}

BigQuery DataFrames

Before trying this sample, follow the BigQuery DataFrames setup instructions in the BigQuery quickstart using BigQuery DataFrames. For more information, see the BigQuery DataFrames reference documentation.

To authenticate to BigQuery, set up Application Default Credentials. For more information, see Set up ADC for a local development environment.

import google.cloud.bigquery

bqclient = google.cloud.bigquery.Client()
bqclient.create_dataset("bqml_tutorial", exists_ok=True)

Import the ONNX model into BigQuery

The following steps show you how to import the sample ONNX model from Cloud Storage by using a CREATE MODEL statement.

To import the ONNX model into your dataset, select one of the following options:

import bigframes
from bigframes.ml.imported import ONNXModel

bigframes.options.bigquery.project = PROJECT_ID
# You can change the location to one of the valid locations: https://cloud.google.com/bigquery/docs/locations#supported_locations
bigframes.options.bigquery.location = "US"

imported_onnx_model = ONNXModel(
    model_path="gs://cloud-samples-data/bigquery/ml/onnx/pipeline_rf.onnx"
)

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

Make predictions with the imported ONNX model

After importing the ONNX model, you use the ML.PREDICT function to make predictions with the model.

The query in the following steps uses imported_onnx_model to make predictions using input data from the iris table in the ml_datasets public dataset. The ONNX model expects four FLOAT values as input:

• sepal_length
• sepal_width
• petal_length
• petal_width

These inputs match the initial_types that were defined when you converted the model into ONNX format.

The outputs include the label and probabilities columns, and the columns from the input table. label represents the predicted class label. probabilities is an array of probabilities representing probabilities for each class.

To make predictions with the imported ONNX model, choose one of the following options:

bq query --use_legacy_sql=false \
'SELECT *
FROM ML.PREDICT(
MODEL `example_dataset.imported_onnx_model`,
(SELECT * FROM `bigquery-public-data.ml_datasets.iris`))'

import bigframes.pandas as bpd

df = bpd.read_gbq("bigquery-public-data.ml_datasets.iris")
predictions = imported_onnx_model.predict(df)
predictions.peek(5)