Function convert_to_torchscript

monai/networks/utils.py:796–863

Utility to convert a model into TorchScript model and save to file, with optional input / output data verification. Args: model: source PyTorch model to save. filename_or_obj: if not None, specify a file-like object (has to implement write and flush) or a st

(
    model: nn.Module,
    filename_or_obj: Any | None = None,
    extra_files: dict | None = None,
    verify: bool = False,
    inputs: Sequence[Any] | None = None,
    device: torch.device | None = None,
    rtol: float = 1e-4,
    atol: float = 0.0,
    use_trace: bool = False,
    **kwargs,
)

def convert_to_torchscript(
    model: nn.Module,
    filename_or_obj: Any | None = None,
    extra_files: dict | None = None,
    verify: bool = False,
    inputs: Sequence[Any] | None = None,
    device: torch.device | None = None,
    rtol: float = 1e-4,
    atol: float = 0.0,
    use_trace: bool = False,
    **kwargs,
):
    """
    Utility to convert a model into TorchScript model and save to file,
    with optional input / output data verification.

    Args:
        model: source PyTorch model to save.
        filename_or_obj: if not None, specify a file-like object (has to implement write and flush)
            or a string containing a file path name to save the TorchScript model.
        extra_files: map from filename to contents which will be stored as part of the save model file.
            for more details: https://pytorch.org/docs/stable/generated/torch.jit.save.html.
        verify: whether to verify the input and output of TorchScript model.
            if `filename_or_obj` is not None, load the saved TorchScript model and verify.
        inputs: input test data to verify model, should be a sequence of data, every item maps to a argument
            of `model()` function.
        device: target device to verify the model, if None, use CUDA if available.
        rtol: the relative tolerance when comparing the outputs of PyTorch model and TorchScript model.
        atol: the absolute tolerance when comparing the outputs of PyTorch model and TorchScript model.
        use_trace: whether to use `torch.jit.trace` to export the TorchScript model.
        kwargs: other arguments except `obj` for `torch.jit.script()` or `torch.jit.trace()` (if use_trace is True)
            to convert model, for more details: https://pytorch.org/docs/master/generated/torch.jit.script.html.

    """
    model.eval()
    with torch.no_grad():
        if use_trace:
            if inputs is None:
                raise ValueError("Missing input data for tracing convert.")
            script_module = torch.jit.trace(model, example_inputs=inputs, **kwargs)
        else:
            script_module = torch.jit.script(model, **kwargs)
        if filename_or_obj is not None:
            torch.jit.save(m=script_module, f=filename_or_obj, _extra_files=extra_files)

    if verify:
        if device is None:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if inputs is None:
            raise ValueError("Missing input data for verification.")

        inputs = [i.to(device) if isinstance(i, torch.Tensor) else i for i in inputs]
        ts_model = torch.jit.load(filename_or_obj) if filename_or_obj is not None else script_module
        ts_model.eval().to(device)
        model = model.to(device)

        with torch.no_grad():
            set_determinism(seed=0)