boundlab.utils.onnx_export

boundlab.utils.onnx_export(f, args, kwargs={}, input_names=None, output_names=None)

Export a PyTorch function to ONNX format.

Examples

>>> import torch
>>> from boundlab.utils import onnx_export
>>> def f(x):
...     return x @ x.T
...
>>> model_proto = onnx_export(f, [3, 4])
>>> list(model_proto.graph)[0].op_type
'MatMul'