PyTorch para ONNX

Este capítulo descreve como converter e exportar modelos PyTorch para modelos ONNX.

Preparação

Configuração do Ambiente

Assim como na etapa de Training, recomendamos que você faça isso em um ambiente virtual durante a fase de exportação do modelo. No ambiente virtual sscma, certifique-se de que a etapa Installation - Prerequisites - Install Extra Dependencies foi concluída.

dica

Se você configurou um ambiente virtual, mas não o ativou, pode ativá-lo com o seguinte comando.

conda activate sscma

Modelos e Pesos

Você também precisa preparar o modelo PyTorch e seus pesos antes de exportar o modelo. Para o modelo, você pode encontrá-lo na seção Config, que já foi pré-configurada. Para os pesos, você pode seguir as etapas abaixo para obter os pesos do modelo.

Consulte a seção Training, escolha um modelo e treine para obter os pesos do modelo.
Ou faça o download dos pesos oficiais pré-treinados do SSCMA no nosso GitHub Releases - Model Zoo.

Exportar Modelo

Para conversão e exportação de modelo, são listados os comandos relevantes com alguns parâmetros comuns.

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

Exemplos de Exportação para ONNX

Aqui estão alguns exemplos de conversão de modelo para referência.

Conversão de Modelo FOMO
Conversão de Modelo PFLD
Conversão de Modelo SWIFT

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'

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'

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'

Validação do Modelo

Como no processo de exportação do modelo o SSCMA fará algumas otimizações no modelo usando algumas ferramentas, como poda de modelo, destilação, etc., embora tenhamos testado e avaliado os pesos do modelo durante o processo de treinamento, recomendamos que você valide novamente o modelo exportado.

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

dica

Para mais parâmetros suportados, consulte o código-fonte tools/inference.py ou execute python3 tools/inference.py --help.

Exemplo de Validação de Modelo

Validação de Modelo FOMO
Validação de Modelo PFLD
Validação de Modelo SWIFT-YOLO

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'

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'

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'

Preparação​

Configuração do Ambiente​

Modelos e Pesos​

Exportar Modelo​

Exemplos de Exportação para ONNX​

Validação do Modelo​

Exemplo de Validação de Modelo​