使用 ONNX 格式的 PyTorch 模型进行预测

概览

开放神经网络交换 (ONNX) 提供了一种用于表示任何机器学习框架的统一格式。BigQuery ML 对 ONNX 的支持让您可以：

• 使用您偏好的框架训练模型。
• 将模型转换为 ONNX 模型格式。如需了解详情，请参阅转换为 ONNX 格式。
• 将 ONNX 模型导入 BigQuery 并使用 BigQuery ML 进行预测。

本教程介绍了如何将使用 PyTorch 训练的 ONNX 模型导入 BigQuery 数据集，并使用这些模型通过 SQL 查询进行预测。您可以使用以下界面导入 ONNX 模型：

• Google Cloud 控制台
• bq 命令行工具中的 bq query 命令
• BigQuery API

如需详细了解如何将 ONNX 模型导入 BigQuery，包括格式和存储要求，请参阅用于导入 ONNX 模型的 CREATE MODEL 语句。

目标

在此教程中，您将学习以下操作：

• 使用 PyTorch 创建和训练模型。
• 使用 torch.onnx 将模型转换为 ONNX 格式。
• 将 ONNX 模型导入 BigQuery 并进行预测。

创建用于图片分类的 PyTorch 视觉模型

导入 PyTorch 预训练的 resnet18，以接受 BigQuery ML ML.DECODE_IMAGE 和 ML.RESIZE_IMAGE 函数返回的解码图片数据。

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()
)

将模型转换为 ONNX 格式并保存

使用 torch.onnx 将 PyTorch 视觉模型导出到名为 resnet18.onnx 的 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

将 ONNX 模型上传到 Cloud Storage

创建一个 Cloud Storage 存储桶来存储 ONNX 模型文件，然后将已保存的 ONNX 模型文件上传到您的 Cloud Storage 存储桶。如需了解详情，请参阅从文件系统上传对象。

将 ONNX 模型导入 BigQuery

此步骤假定您已将 ONNX 模型上传到 Cloud Storage 存储桶。示例模型存储在 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/*')"
}

在 BigQuery 中创建对象表以访问图片数据

如需访问 BigQuery 中的非结构化数据，您需要创建一个对象表。如需查看详细说明，请参阅创建对象表。

在存储于 gs://mybucket/goldfish.jpg 的金鱼图片上创建名为 goldfish_image_table 的对象表。

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');

使用导入的 ONNX 模型进行预测

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))"
}

后续步骤

• 如需详细了解如何导入 ONNX 模型，请参阅用于 ONNX 模型的 CREATE MODEL 语句。
• 如需详细了解可用的 ONNX 转换器和教程，请参阅转换为 ONNX 格式。
• 如需大致了解 BigQuery ML，请参阅 BigQuery ML 简介。
• 如需开始使用 BigQuery ML，请参阅在 BigQuery ML 中创建机器学习模型。
• 如需详细了解如何使用模型，请参阅以下资源：
  • 获取模型元数据
  • 更新模型元数据
  • 管理模型
  • 删除模型