详细信息可以看官方文档:./docs/en/tutorials/pytorch2onnx.md
这里把命令摘了出来
用法
//bash
python tools/deployment/pytorch2onnx.py \
${CONFIG_FILE} \
${CHECKPOINT_FILE} \
--output-file ${OUTPUT_FILE} \
--input-img ${INPUT_IMAGE_PATH} \
--shape ${IMAGE_SHAPE} \
--test-img ${TEST_IMAGE_PATH} \
--opset-version ${OPSET_VERSION} \
--cfg-options ${CFG_OPTIONS}
--dynamic-export \
--show \
--verify \
--simplify \
Description of all arguments
config
: The path of a model config file.checkpoint
: The path of a model checkpoint file.--output-file
: The path of output ONNX model. If not specified, it will be set totmp.onnx
.--input-img
: The path of an input image for tracing and conversion. By default, it will be set totests/data/color.jpg
.--shape
: The height and width of input tensor to the model. If not specified, it will be set to800 1216
.--test-img
: The path of an image to verify the exported ONNX model. By default, it will be set toNone
, meaning it will use--input-img
for verification.--opset-version
: The opset version of ONNX. If not specified, it will be set to11
.--dynamic-export
: Determines whether to export ONNX model with dynamic input and output shapes. If not specified, it will be set toFalse
.--show
: Determines whether to print the architecture of the exported model and whether to show detection outputs when--verify
is set toTrue
. If not specified, it will be set toFalse
.--verify
: Determines whether to verify the correctness of an exported model. If not specified, it will be set toFalse
.--simplify
: Determines whether to simplify the exported ONNX model. If not specified, it will be set toFalse
.--cfg-options
: Override some settings in the used config file, the key-value pair inxxx=yyy
format will be merged into config file.--skip-postprocess
: Determines whether export model without post process. If not specified, it will be set toFalse
. Notice: This is an experimental option. Only work for some single stage models. Users need to implement the post-process by themselves. We do not guarantee the correctness of the exported model.
-config
:模型配置文件的路径。
-“checkpoint”:模型检查点文件的路径。
-“–output file”:NX模型的输出路径。如果未指定,它将被设置为“tmp”。onnx`。
-“–input img”:用于跟踪和转换的输入图像的路径。默认情况下,它将设置为“测试/数据/颜色”。jpg`。
-“–shape”:模型输入张量的高度和宽度。如果未指定,它将设置为“800 1216”。
-“–test img”:验证导出的ONNX模型的图像路径。默认情况下,它将被设置为“无”,这意味着它将使用“–input img”进行验证。
-“–opset version”:ONNX的opset版本。如果未指定,它将被设置为“11”。
-“–dynamic export”:确定是否导出带有动态输入和输出形状的ONNX模型。如果未指定,它将被设置为“False”。
-“–show”:确定当“–verify”设置为“True”时,是否打印导出模型的体系结构,以及是否显示检测输出。如果未指定,它将被设置为“False”。
-“–verify”:确定是否验证导出模型的正确性。如果未指定,它将被设置为“False”。
-“–simplify”:确定是否简化导出的ONNX模型。如果未指定,它将被设置为“False”。
-“–cfg options:覆盖所用配置文件中的一些设置,
xxx=yyy`格式的键值对将合并到配置文件中。
-“–skip postprocess”:确定是否在不进行后期处理的情况下导出模型。如果未指定,它将被设置为“False”。注意:这是一个实验选项。仅适用于某些单级模型。用户需要自己实现post流程。我们不保证导出模型的正确性。
使用例子
//bash
python tools/deployment/pytorch2onnx.py \
configs/yolo/yolov3_d53_mstrain-608_273e_coco.py \
checkpoints/yolo/yolov3_d53_mstrain-608_273e_coco.pth \
--output-file checkpoints/yolo/yolov3_d53_mstrain-608_273e_coco.onnx \
--input-img demo/demo.jpg \
--test-img tests/data/color.jpg \
--shape 608 608 \
--show \
--verify \
--dynamic-export \
--cfg-options \
model.test_cfg.deploy_nms_pre=-1 \
过程中需要安装一些包,缺啥装啥就完事了
pip install opencv-python
pip install onnx
pip install onnxruntime
转完onnx模型可以使用Netron可视化网络结构
网址:https://netron.app/