1. 环境准备

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模型权重文件

更多推荐