ONNX¶

关于此页面

这是 BentoML 中 ONNX 的 API 参考。有关如何在 BentoML 中使用 ONNX 的更多信息，请参阅 ONNX 指南。

bentoml.onnx.save_model(name: Tag | str, model: onnx.ModelProto, *, signatures: dict[str, ModelSignatureDict] | dict[str, ModelSignature] | None = None, labels: dict[str, str] | None = None, custom_objects: dict[str, t.Any] | None = None, external_modules: t.List[ModuleType] | None = None, metadata: dict[str, t.Any] | None = None) → bentoml.Model¶

将 ONNX 模型实例保存到 BentoML 模型库。

参数:

name (str) – 在 BentoML 模型库中为模型指定的名称。这必须是有效的 Tag 名称。
model (ModelProto) – 要保存的 ONNX 模型。
signatures (dict[str, ModelSignatureDict], 可选) – 要使用的预测方法的签名。如果未提供，签名默认为 {"run": {"batchable": False}}。有关更多详细信息，请参阅 ModelSignature。bentoml.onnx 内部使用 onnxruntime.InferenceSession 运行推理。当原始模型转换为 ONNX 格式并由 onnxruntime.InferenceSession 加载时，原始模型的推理方法被转换为 onnxruntime.InferenceSession 的 run 方法。signatures 在此指代 onnxruntime.InferenceSession 的预测方法，因此 signatures 中唯一允许的方法名称是 run。
labels (dict[str, str], 可选) – 与模型关联的默认管理标签集。例如 {"training-set": "data-1"}。
custom_objects (dict[str, Any], 可选) –
要随模型保存的自定义对象。例如 {"my-normalizer": normalizer}。

自定义对象目前使用 cloudpickle 序列化，但此实现可能会更改。
external_modules (List[ModuleType], 可选, 默认为 None) – 用户定义的额外 Python 模块，随模型或自定义对象一起保存，例如分词器模块、预处理器模块、模型配置模块
metadata (dict[str, Any], 可选) –
与模型关联的元数据。例如 {"bias": 4}。

元数据用于在模型管理 UI 中显示，因此必须是默认的 Python 类型，例如 str 或 int。

返回:

包含已保存 ONNX 模型实例的 BentoML 模型。

返回类型:

Model

示例

import bentoml

import torch
import torch.nn as nn

class ExtendedModel(nn.Module):
    def __init__(self, D_in, H, D_out):
        # In the constructor we instantiate two nn.Linear modules and assign them as
        #  member variables.
        super(ExtendedModel, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)

    def forward(self, x, bias):
        # In the forward function we accept a Tensor of input data and an optional bias
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred + bias


N, D_in, H, D_out = 64, 1000, 100, 1
x = torch.randn(N, D_in)
model = ExtendedModel(D_in, H, D_out)

input_names = ["x", "bias"]
output_names = ["output1"]

tmpdir = "/tmp/model"
model_path = os.path.join(tmpdir, "test_torch.onnx")
torch.onnx.export(
    model,
    (x, torch.Tensor([1.0])),
    model_path,
    input_names=input_names,
    output_names=output_names,
)

bento_model = bentoml.onnx.save_model("onnx_model", model_path, signatures={"run": {"batchable": True}})

bentoml.onnx.load_model(bento_model: str | Tag | bentoml.Model, *, providers: ProvidersType | None = None, session_options: ort.SessionOptions | None = None) → ort.InferenceSession¶

从本地 BentoML 模型库加载具有给定标签的 ONNX 模型。

参数:

bento_model (str | Tag | Model) – 要从模型库获取的模型标签，或用于从中加载模型的 BentoML ~bentoml.Model 实例。
providers (List[Union[str, Tuple[str, Dict[str, Any]], 可选, 默认为 None) – 用户提供的不同提供者。默认情况下，BentoML 在加载模型时将使用 ["CPUExecutionProvider"]。
session_options (onnxruntime.SessionOptions, 可选, 默认为 None) – 每个用例的 SessionOptions。如果未指定，则默认为 None。

返回:

使用从模型库加载的 ONNX 模型创建的 ONNX Runtime 推理会话实例。

返回类型:

onnxruntime.InferenceSession

示例

import bentoml
sess = bentoml.onnx.load_model("my_onnx_model")

bentoml.onnx.get(tag_like: str | Tag) → Model¶

获取具有给定标签的 BentoML 模型。

参数:: tag_like – 要从模型库检索的模型标签。
返回:: 具有匹配标签的 BentoML Model。
返回类型:: Model

示例

import bentoml
# target model must be from the BentoML model store
model = bentoml.onnx.get("onnx_resnet50")