peal.architectures.onnx_predictor.export_to_onnx

peal.architectures.onnx_predictor.export_to_onnx(model, path, input_shape=None, device=None, opset_version=None)[source]

Export a predictor to ONNX with a dynamic batch axis.

input_shape is [C, H, W] (or the full [B, C, H, W]; the batch axis is replaced). An OnnxPredictor supplies its own shape when none is given. Returns path.

Parameters:

  • model (torch.nn.Module) – Predictor to export; switched to eval mode.

  • path (str) – Destination .onnx file.

  • input_shape (list of int, optional) – See above; symbolic (None) dims are set to 1.

  • device (torch.device, optional) – Device of the dummy input; defaults to the model’s first parameter (CPU for parameterless models).

  • opset_version (int, optional) – Passed to torch.onnx.export when given.

Returns:

path.

Return type:

str

Raises:

ValueError – If no input_shape is available.

Notes

Uses the legacy (dynamo=False) exporter with input "input", output "output" and a dynamic batch_size axis on both.