aimet_torch.onnx.export

aimet_torch.onnx.export(model, args, f, *, export_int32_bias=True, prequantize_constants=False, force_activation_as='unsigned', **kwargs)

Export QuantizationSimModel to onnx model with onnx QuantizeLinear and DequantizeLinear embedded in the graph.

This function takes the same set of arguments as torch.onnx.export()

Parameters:

• model (Union[Module, QuantizationSimModel]) – The model to be exported

• args (Union[Tuple[Any, ...], Tensor]) – Same as torch.onnx.export()

• f (Union[str, BytesIO]) – Same as torch.onnx.export()

• export_int32_bias (bool) – If true, generate and export int32 bias encoding on the fly (default: True).

• prequantize_constants (bool) – If True, quantized weights will be represented as quantized weight followed by DequantizeLinear. If False, they will be represented as floating point weight followed by QuantizeLinear and DequantizeLinear (default: False).

• force_activation_as (Union[Literal['unsigned'], Literal['signed'], None]) – Force representing quantized activations as signed or unsigned integers (default: “unsigned”)

• **kwargs – Same as torch.onnx.export()

Note

For robustness, onnx >=1.19 is highly recommended with this API, especially when exporting large models (>2GB). This is due to a known bug in onnx <1.19 version converter. For more information, see https://github.com/onnx/onnx/issues/6529

Examples

>>> aimet_torch.onnx.export(sim.model, x, f="model.onnx",
...                         input_names=["input"], output_names=["output"],
...                         opset_version=21, dynamo=False,
...                         export_int32_bias=True)
...
>>> import onnxruntime as ort
>>> options = ort.SessionOptions()
>>> options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
>>> sess = ort.InferenceSession("model.onnx", sess_options=options)
>>> onnx_output, = sess.run(None, {"input": x.detach().numpy()})
>>> torch.nn.functional.cosine_similarity(torch.from_numpy(onnx_output), sim.model(x))
tensor([1.0000, 0.9999, 1.0000,  ..., 1.0000, 1.0000, 1.0000],
       grad_fn=<AliasBackward0>)