Export trained PyTorch flow model to TorchScript model

This Python script allows you to export a trained PyTorch flow model to a TorchScript model that can be easily deployed in various runtime environment, 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_flow when installing Metavision SDK from installer or packages. For other deployment methods, check the page Path of Samples.

Expected Output

Compiled TorchScript Model that can be easily deployed during runtime.

Specifically, it will output:

  1. model.ptjit (the model).

  2. info_ssd_jit.json (the hyperparameters used during training)

Setup & requirements

You will need to provide the following input:

  • path to the checkpoint. You can use flow_model_alpha.ckpt from our pre-trained models

  • path to the output directory

How to start

To run the script with flow_model_alpha.ckpt:

python export_flow.py flow_model_alpha.ckpt /path/to/output

To find the full list of options, run:

python export_flow.py -h