ICode9

精准搜索请尝试: 精确搜索
首页 > 其他分享> 文章详细

Pytorch模型转onnx,变得很大

2021-09-12 15:04:18  阅读:416  来源: 互联网

标签:onnx 模型 same initializer raw Pytorch graph model


引言

  • 自己用PyTorch训练一个模型,pth模型大小有6.58M
  • 寻思着转换为onnx格式,便于推理和部署
  • 不料转换后模型竟然增加到242M
模型名称大小
raw.pth6.58M
convert_raw.onnx242M

解决方案

  • 可能原因:onnx中有大量算子重复
  • 解决代码 [来源]
from onnxruntime.transformers.onnx_model import OnnxModel
import onnx

def has_same_value(val_one,val_two):
    if val_one.raw_data == val_two.raw_data:
        return True
    else:
        return False

path = f"convert_raw.onnx"  # 242M
output_path = f"slim_convert.onnx"  # 7.50M
model = onnx.load(path)
onnx_model = OnnxModel(model)

count = len(model.graph.initializer)
same = [-1] * count
for i in tqdm(range(count - 1)):
  if same[i] >= 0:
        continue
  for j in range(i+1, count):
      if has_same_value(model.graph.initializer[i], 
                        model.graph.initializer[j]):
          same[j] = i

for i in tqdm(range(count)):
   if same[i] >= 0:
       onnx_model.replace_input_of_all_nodes(model.graph.initializer[i].name,
                                             model.graph.initializer[same[i]].name)
onnx_model.update_graph()
onnx_model.save_model_to_file(output_path)

相关资料

标签:onnx,模型,same,initializer,raw,Pytorch,graph,model
来源: https://blog.csdn.net/shiwanghualuo/article/details/120250602

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有