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
如果你需要其它模型,请修改它。已经存在的模型可以在第一个注释掉的网址中找到。