Enterprise-ready security: ISO 27001 + SOC 2 Type I compliant.

Link to this sectionReference for ultralytics/utils/export/torchscript.py#

Improvements

This page is sourced from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/torchscript.py. Have an improvement or example to add? Open a Pull Request — thank you! 🙏


Summary

Link to this sectionFunction ultralytics.utils.export.torchscript.torch2torchscript#

def torch2torchscript(
    model: torch.nn.Module,
    im: torch.Tensor,
    output_file: Path | str,
    metadata: dict | None = None,
    prefix: str = "",
) -> str

Export a PyTorch model to TorchScript format.

Args

NameTypeDescriptionDefault
modeltorch.nn.ModuleThe PyTorch model to export (may be NMS-wrapped).required
imtorch.TensorExample input tensor for tracing.required
output_filePath | strPath to save the exported TorchScript model.required
metadatadict | NoneOptional metadata to embed in the TorchScript archive.None
prefixstrPrefix for log messages.""

Returns

TypeDescription
strPath to the exported .torchscript file.
GitHubultralytics/utils/export/torchscript.py
def torch2torchscript(
    model: torch.nn.Module,
    im: torch.Tensor,
    output_file: Path | str,
    metadata: dict | None = None,
    prefix: str = "",
) -> str:
    """Export a PyTorch model to TorchScript format.

    Args:
        model (torch.nn.Module): The PyTorch model to export (may be NMS-wrapped).
        im (torch.Tensor): Example input tensor for tracing.
        output_file (Path | str): Path to save the exported TorchScript model.
        metadata (dict | None): Optional metadata to embed in the TorchScript archive.
        prefix (str): Prefix for log messages.

    Returns:
        (str): Path to the exported ``.torchscript`` file.
    """
    LOGGER.info(f"\n{prefix} starting export with torch {TORCH_VERSION}...")

    output_file = str(output_file)
    ts = torch.jit.trace(model, im, strict=False)
    extra_files = {"config.txt": json.dumps(metadata or {})}  # torch._C.ExtraFilesMap()
    ts.save(output_file, _extra_files=extra_files)
    return output_file