Implementation:Bentoml BentoML Framework TorchScript
| Knowledge Sources | |
|---|---|
| Domains | ML Framework, TorchScript, Deep Learning, Model Serialization |
| Last Updated | 2026-02-13 15:00 GMT |
Overview
Provides the BentoML framework integration for TorchScript, enabling saving, loading, and serving torch.ScriptModule models through the BentoML model store.
Description
This module implements the BentoML framework adapter for TorchScript (torch.jit). It uses torch.jit.save() and torch.jit.load() for serialization, which produces optimized, portable models that can run without Python dependencies.
save_model() validates the input is a torch.ScriptModule (or torch.jit.ScriptModule), creates a ModelContext with framework version information, and saves the model to saved_model.pt using torch.jit.save(). It supports an _extra_files parameter for bundling additional files with the model (recorded in metadata). When called with _framework_name="pytorch_lightning", it records both torch and pytorch_lightning versions (used by the PyTorch Lightning adapter). The default signature is __call__ (non-batchable). It uses PartialKwargsModelOptions for method-level default arguments.
load_model() loads the model with torch.jit.load(), supporting a device_id parameter for device placement and an _extra_files parameter for retrieving bundled extra files. It can return either a torch.ScriptModule or a tuple of (ScriptModule, extra_files_dict) depending on whether _extra_files is provided.
get() retrieves the BentoML Model metadata for a given tag.
get_runnable() creates a PytorchModelRunnable using shared PyTorch utilities (partial_class, make_pytorch_runnable_method), supporting partial kwargs per method.
Usage
Use this module to save TorchScript models (created via torch.jit.script() or torch.jit.trace()) to the BentoML model store and serve them. Access via bentoml.torchscript. Also used internally by bentoml.pytorch_lightning.
Code Reference
Source Location
- Repository: Bentoml_BentoML
- File: src/bentoml/_internal/frameworks/torchscript.py
- Lines: 1-196
Signature
def get(tag_like: str | Tag) -> Model: ...
def load_model(
bentoml_model: str | Tag | Model,
device_id: str | None = "cpu",
*,
_extra_files: dict[str, t.Any] | None = None,
) -> torch.ScriptModule | tuple[torch.ScriptModule, dict[str, t.Any]]: ...
def save_model(
name: Tag | str,
model: torch.ScriptModule,
*,
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,
_framework_name: str = "torchscript",
_module_name: str = MODULE_NAME,
_extra_files: dict[str, t.Any] | None = None,
) -> bentoml.Model: ...
def get_runnable(bento_model: Model) -> type[PytorchModelRunnable]: ...
Import
import bentoml
# Via the public API
bento_model = bentoml.torchscript.save_model("scripted_model", script_module)
loaded = bentoml.torchscript.load_model("scripted_model:latest", device_id="cuda:0")
# Direct import
from bentoml._internal.frameworks.torchscript import save_model, load_model, get, get_runnable
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| name | Tag or str | Yes (save_model) | Name or tag for the model in the store |
| model | torch.ScriptModule | Yes (save_model) | The TorchScript module to save |
| tag_like | str or Tag | Yes (get) | Tag to retrieve from the model store |
| bentoml_model | str, Tag, or Model | Yes (load_model) | Model identifier or object to load |
| device_id | str or None | No (default "cpu") | Device to load the model onto (e.g., "cpu", "cuda:0") |
| _extra_files | dict[str, Any] or None | No | Additional files to bundle with or load from the model |
| signatures | ModelSignaturesType or None | No | Method signatures; defaults to {"__call__": {"batchable": False}} |
| _framework_name | str | No (default "torchscript") | Internal: framework name for context; set to "pytorch_lightning" by the PL adapter |
| _module_name | str | No (default MODULE_NAME) | Internal: module name for the model info |
| labels | dict[str, str] or None | No | User-defined labels for model management |
| custom_objects | dict[str, Any] or None | No | Additional objects to save with the model |
| external_modules | List[ModuleType] or None | No | Additional Python modules to bundle |
| metadata | dict[str, Any] or None | No | Custom metadata for the model |
Outputs
| Name | Type | Description |
|---|---|---|
| bentoml.Model (save_model) | bentoml.Model | The saved model reference in the BentoML store |
| torch.ScriptModule (load_model) | torch.ScriptModule or tuple | The deserialized TorchScript module; tuple if _extra_files is provided |
| Model (get) | Model | BentoML Model metadata object |
| PytorchModelRunnable (get_runnable) | type[Runnable] | A Runnable class using shared PyTorch utilities with partial kwargs support |
Usage Examples
import torch
import bentoml
# Create a TorchScript model via scripting
class MyModel(torch.nn.Module):
def forward(self, x):
return x * 2
scripted = torch.jit.script(MyModel())
# Save to BentoML
tag = bentoml.torchscript.save_model("my_scripted_model", scripted)
print(f"Saved: {tag}")
# Load on a specific device
model = bentoml.torchscript.load_model("my_scripted_model:latest", device_id="cuda:0")
# Save with extra files
extra = {"config.json": "{}"}
tag = bentoml.torchscript.save_model(
"model_with_config", scripted, _extra_files=extra
)
# Load with extra files
extra_out = {"config.json": ""}
model = bentoml.torchscript.load_model(
"model_with_config:latest", _extra_files=extra_out
)
print(extra_out["config.json"])
# Create a TorchScript model via tracing
traced = torch.jit.trace(MyModel(), torch.randn(1, 10))
tag = bentoml.torchscript.save_model("traced_model", traced)