PyTorch

关于此页

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

注意

您可以在我们的 BentoML/examples 目录下找到更多 PyTorch 的示例。

bentoml.pytorch.save_model(name: Tag | str, model: torch.nn.Module, *, signatures: ModelSignaturesType | None = None, labels: t.Dict[str, str] | None = None, custom_objects: t.Dict[str, t.Any] | None = None, external_modules: t.List[ModuleType] | None = None, metadata: t.Dict[str, t.Any] | None = None) → bentoml.Model

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

参数：

• name (str) – 给定模型实例的名称。这应该通过 Python 标识符检查。

• model (torch.nn.Module) – 要保存的模型实例

• signatures (ModelSignaturesType, 可选, 默认为 None) – 方法名称及其对应签名的字典。

• labels (Dict[str, str], 可选, 默认为 None) – 用于管理模型的用户定义标签，例如 team=nlp, stage=dev

• custom_objects (Dict[str, Any]], 可选, 默认为 None) – 要与模型一起保存的用户定义的附加 Python 对象，例如 tokenizer 实例、预处理器函数、模型配置 json

• external_modules (List[ModuleType], 可选, 默认为 None) – 要与模型或自定义对象一起保存的用户定义的附加 Python 模块，例如 tokenizer 模块、预处理器模块、模型配置模块

• metadata (Dict[str, Any], 可选, 默认为 None) – 给定模型的自定义元数据。

返回：

一个 tag，格式为 name:version，其中 name 是用户定义的模型名称，version 是由 BentoML 生成的版本。

返回类型：

Tag

示例

import torch
import bentoml

class NGramLanguageModeler(nn.Module):

    def __init__(self, vocab_size, embedding_dim, context_size):
        super(NGramLanguageModeler, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.linear1 = nn.Linear(context_size * embedding_dim, 128)
        self.linear2 = nn.Linear(128, vocab_size)

    def forward(self, inputs):
        embeds = self.embeddings(inputs).view((1, -1))
        out = F.relu(self.linear1(embeds))
        out = self.linear2(out)
        log_probs = F.log_softmax(out, dim=1)
        return log_probs

tag = bentoml.pytorch.save("ngrams", NGramLanguageModeler(len(vocab), EMBEDDING_DIM, CONTEXT_SIZE))
# example tag: ngrams:20201012_DE43A2

与 Torch Hub 和 BentoML 集成

import torch
import bentoml

resnet50 = torch.hub.load("pytorch/vision", "resnet50", pretrained=True)
...
# trained a custom resnet50

tag = bentoml.pytorch.save("resnet50", resnet50)

bentoml.pytorch.load_model(bentoml_model: str | Tag | Model, device_id: str | None = 'cpu', **torch_load_args: Any) → torch.nn.Module

从具有给定名称的 BentoML 模型加载模型。

参数：

• tag (Union[str, Tag]) – BentoML 本地模型库中已保存模型的 Tag。

• device_id (str, 可选, 默认为 cpu) – 可选设备，用于放置给定的模型。请参阅 设备属性。

返回：

一个 torch.nn.Module 实例，来自 BentoML 模型库。

返回类型：

torch.nn.Module

示例

import bentoml
model = bentoml.pytorch.load_model('lit_classifier:latest', device_id="cuda:0")

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