PyTorch to ONNX

This chapter will describe how to convert and export PyTorch models to ONNX models.

Preparation

Environment Configuration

As the Training step, we recommend you to do it in a virtual environment during the model exporting phase. In the sscma virtual environment, make sure that the Installation - Prerequisites - Install Extra Dependencies step has been completed.

TIP

If you have configured a virtual environment but not activated it, you can activate it with the following command.

conda activate sscma

Models and Weights

You also need to prepare the PyTorch model and its weights before exporting the model. For the model, you can find it in the Config section, we have already preconfigured. For the weights, you can refer to the following steps to get the model weights.

• Refer to Training section and choose a model, and train to get the model weights.

• Or download the SSCMA official pre-trained weights from our GitHub Releases - Model Zoo.

Export Model

For model convert and export, the relevant commands with some common parameters are listed.

python3 tools/export.py \
    "<CONFIG_FILE_PATH>" \
    "<CHECKPOINT_FILE_PATH>" \
    --target onnx

ONNX Export Examples

Here are some model conversion examples for reference.

FOMO Model Conversion

python3 tools/export.py \
    configs/fomo/fomo_mobnetv2_0.35_x8_abl_coco.py \
    "$(cat work_dirs/fomo_mobnetv2_0.35_x8_abl_coco/last_checkpoint)" \
    --target onnx \
    --cfg-options \
        data_root='datasets/mask'

PFLD Model Conversion

python3 tools/export.py \
    configs/pfld/pfld_mbv2n_112.py \
    "$(cat work_dirs/pfld_mbv2n_112/last_checkpoint)" \
    --target onnx \
    --cfg-options \
        data_root='datasets/meter'

SWIFT-YOLO Model Conversion

python3 tools/export.py \
    configs/swift_yolo/swift_yolo_tiny_1xb16_300e_coco.py \
    "$(cat work_dirs/swift_yolo_tiny_1xb16_300e_coco/last_checkpoint)" \
    --target onnx \
    --cfg-options \
        data_root='datasets/digital_meter'

Model Validation

Since in the process of exporting the model, SSCMA will do some optimization for the model using some tools, such as model pruning, distillation, etc. Although we have tested and evaluated the model weights during the training process, we recommend you to validate the exported model again.

python3 tools/inference.py \
    "<CONFIG_FILE_PATH>" \
    "<CHECKPOINT_FILE_PATH>" \
    --show \
    --cfg-options "<CFG_OPTIONS>"

TIP

For more parameters supported, please refer to the source code tools/inference.py or run python3 tools/inference.py --help.

Model Validation Example

FOMO Model Validation

python3 tools/inference.py \
    configs/fomo/fomo_mobnetv2_0.35_x8_abl_coco.py \
    "$(cat work_dirs/fomo_mobnetv2_0.35_x8_abl_coco/last_checkpoint | sed -e 's/.pth/.onnx/g')" \
    --show \
    --cfg-options \
        data_root='datasets/mask'

PFLD Model Validation

python3 tools/inference.py \
    configs/pfld/pfld_mbv2n_112.py \
    "$(cat work_dirs/pfld_mbv2n_112/last_checkpoint | sed -e 's/.pth/.onnx/g')" \
    --show \
    --cfg-options \
        data_root='datasets/meter'

SWIFT-YOLO Model Validation

python3 tools/inference.py \
    configs/swift_yolo/swift_yolo_tiny_1xb16_300e_coco.py \
    "$(cat work_dirs/swift_yolo_tiny_1xb16_300e_coco/last_checkpoint | sed -e 's/.pth/.onnx/g')" \
    --show \
    --cfg-options \
        data_root='datasets/digital_meter'