pytorch中一些CV模型的预训练模型下载并转换为onnx模型

import torch
import onnx

# https://github.com/rwightman/gen-efficientnet-pytorch/blob/master/hubconf.py
model=torch.hub.load("rwightman/pytorch-image-models","mobilenetv3_large_100")
model.eval()
batch_size=1
input_shape = (3,244,244)
export_onnx_file="mobilenetv3_large_100.onnx"
input_vars = torch.randn(batch_size,*input_shape)
model.eval()

#torch.onnx.export(model, args=(input_var), f="cnn.onnx", opset_version=12,verbose=False, input_names=["input"], output_names=["output"])
torch.onnx.export(model,
                    input_vars,
                    export_onnx_file,
                    opset_version=13,
                    do_constant_folding=True,	# 是否执行常量折叠优化
                    input_names=["input"],		# 输入名
                    output_names=["output"],	# 输出名
                    dynamic_axes={"input":{0:"batch_size"},		# 批处理变量
                                    "output":{0:"batch_size"}})


onnx_model = onnx.load(export_onnx_file)
onnx.checker.check_model(onnx_model)
    

注意上面的代码中的模型名字：mobilenetv3_large_100

如果你需要其它模型，请修改它。已经存在的模型可以在第一个注释掉的网址中找到。