Function export_segmentation_onnx

samples/common.py:300–345

Export an ONNX model from a TorchVision segmentation model. Args: model: PyTorch model to export. output: Path to the output ONNX model. input_shape: Shape of the input tensor. verbose: Whether to print the verbose output.

(
    model: "torch.nn.Module",  # noqa: F821
    output: Path,
    input_shape: tuple[int, ...],
    *,
    verbose: bool | None = None,
)

def export_segmentation_onnx(
    model: "torch.nn.Module",  # noqa: F821
    output: Path,
    input_shape: tuple[int, ...],
    *,
    verbose: bool | None = None,
) -> None:
    """
    Export an ONNX model from a TorchVision segmentation model.

    Args:
        model: PyTorch model to export.
        output: Path to the output ONNX model.
        input_shape: Shape of the input tensor.
        verbose: Whether to print the verbose output.
    """
    import torch
    import onnx
    import onnxslim

    class SegmentationEnd2End(torch.nn.Module):
        def __init__(self, model):
            super().__init__()
            self.model = model

        def forward(self, x: torch.Tensor) -> torch.Tensor:
            softmax = torch.nn.functional.softmax(self.model(x)["out"], dim=1)
            return softmax

    model = SegmentationEnd2End(model)
    model.eval()

    torch.onnx.export(
        model,
        args=(torch.randn(1, *input_shape),),
        f=output,
        export_params=True,
        opset_version=13,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        verbose=verbose,
    )
    onnx_model = onnx.load(output)
    slim_onnx_model = onnxslim.slim(onnx_model)
    onnx.save(slim_onnx_model, output)


def export_retinanet_onnx(