标签：脚本 name ONNX 模型 Optimization output input Model model
简介
我们都知道目前流行的深度学习框架有很多种，例如Pytorch、TensorFlow、Caffe等，在实际训练和部署的过程中，常常会接触到不同框架模型之间的转换，大多数时候这些并不复杂，只需要写一些简短的脚本就能够完成。而像OpenVINO中IR这样的特殊模型，我们也可以通过其提供的自带脚本进行转换，下面简单介绍常用的模型之间转换的脚本和方法。
Pth模型 转 ONNX模型
import torch
import argparse
parser = argparse.ArgumentParser()
# the function about making the ".pth" model trained by pytorch convert to the "ONNX" model
def convert_model_to_ONNX(input_pth_model, output_ONNX_model, model_name):
dummy_input = torch.randn(1, 1, 64, 64)
model = torch.load(input_pth_model, map_location='cpu')
input_names = ["input"]
output_names = ["output"]
model.eval()
torch.onnx.export(model, dummy_input, output_ONNX_model + '/' + model_name, verbose=True, input_names=input_names,output_names=output_names)
if __name__ == '__main__':
parser.add_argument("--input", type=str, help='input pth model path')
parser.add_argument("--output", type=str, help='output onnx model path')
parser.add_argument("--name", type=str, help='output onnx model filename')
opt = parser.parse_args()
# call the function named "convert_model_to_ONNX".
convert_model_to_ONNX(opt.input, opt.output, opt.name)
Pth模型 转 Pt模型
import torch
from collections import OrderedDict
import pickle
import os
import sys
sys.path.insert(0, 'model dir')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if __name__ == '__main__':
modelfile = 'yolov5s.pt'
utl_model = torch.load('your output dir/'+modelfile)
utl_param = utl_model['model'].model
torch.save(utl_param.state_dict(), os.path.splitext(
modelfile)[0] + '_param.pth')
own_state = utl_param.state_dict()
print(len(own_state))
numpy_param = OrderedDict()
for name in own_state:
numpy_param[name] = own_state[name].data.cpu().numpy()
print(len(numpy_param))
with open(os.path.splitext(modelfile)[0] + '_numpy_param.pkl', 'wb') as fw:
pickle.dump(numpy_param, fw)
Pt模型 转 ONNX模型
使用YOLOv5源码中的export.py脚本，可以很容易地将yolo训练出的pt模型转化成通用的onnx模型。
在yolov5的models文件夹下
python3 export.py --weights 模型路径 --img-size 输入层尺寸
即可直接将pt模型转化成onnx模型。
ONNX模型 转 IR模型
OpenVINO推理加速使用的是一种中间表示(IR)模型，其自带的模型优化器中也提供了ONNX模型转换成IR模型的脚本，在deployment_tools/model_optimizer路径下
运行模型转换脚本
cd <INSTALL_DIR>/deployment_tools/model_optimizer/
python3 mo.py --input_model <INPUT_MODEL>.onnx --output_dir <OUTPUT_MODEL_DIR> --model_name <OUTPUT_MODEL_NAME>
出现SUCCESS则表示转换成功，在对应的路径下会生成.xml和.bin文件，在任何时候，你都需要将这两个模型文件放在统一路径下。
参考
Converting a Model to Intermediate Representation (IR)
后续
喜欢的话可以关注一下我的公众号技术开发小圈，尤其是对深度学习以及计算机视觉有兴趣的朋友，我会把相关的源码以及更多资料发在上面，希望可以帮助到新入门的大家！
标签：脚本,name,ONNX,模型,Optimization,output,input,Model,model 来源： https://blog.csdn.net/Rosen_er/article/details/123589424
本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享； 2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关； 3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关； 4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除； 5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。