Export trained PyTorch classification model to ONNX model

This Python script allows you to export a trained PyTorch classification model to an ONNX model that can be easily deployed in various runtime environments, with an optimized latency and throughput.

The source code of this sample can be found in <install-prefix>/share/metavision/sdk/ml/python_samples/export_fnn_classifier_onnx.py when installing Metavision SDK from installer or packages. For other deployment methods, check the page Path of Samples.

Expected Output

The classification ONNX Model that can be easily deployed during runtime.

Specifically, it will produce a model_classifier.onnx file.

Setup & requirements

You will need to provide the following input:

  • path to the checkpoint. For example, you can use chifoumi_fnn.ckpt from our pre-trained models.

  • path to the output directory

How to start

To run the script:

python export_fnn_classifier_onnx.py classifier_model.ckpt <path_to_output>

Note

  1. use –use_onnx_simplifier to simplify the onnx model with onnxsim library before exporting.

  2. use –use_dynamic_axes to make the dimension of batch_size` as the dynamic axis.

To find the full list of options, run:

python export_fnn_classifier_onnx.py -h