pytorch模型转onnx、openvino格式最佳实践

1. 环境准备

• conda环境：

  conda create -n py310 python=3.10
  conda activate py310

• torch依赖：

  conda install pytorch torchvision torchaudio cpuonly -c pytorch

• openvino工具：

  pip install openvino-dev[extras]

  不能仅安装openvino，模型转换时需要model optimizer工具，因此需要使用开发者版本。extras选填 pytorch,onnx，其他选择参考 https://docs.openvino.ai/2024/documentation/legacy-features/install-dev-tools.html#install-dev-tools

2. 模型转换（以resnet50为例）

• 加载pytorch模型

import torch

import torchvision.models as models

#加载预训练的ResNet50模型

model = models.resnet50(pretrained=True)

#保存为pt文件 torch.save(model.state_dict(), 'resnet50.pt')

#加载模型并设为评估模式

model.load_state_dict(torch.load('resnet50.pt'))

model.eval()

#创建一个示例输入，导出为ONNX格式模型

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(model, dummy_input, 'resnet50.onnx', opset_version=11)

• 将ONNX模型转换为openvino模型

mo --input_model resnet50.onnx --output_dir ./openvino_model

• 检查输出

生成resnet50.xml模型结构文件和resnet50.bin模型权重文件