Pytorch模型转onnx，变得很大

引言

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

模型名称	大小
raw.pth	6.58M
convert_raw.onnx	242M
slim_convert_raw.onnx	7.5M

解决方案

• 可能原因：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)

相关资料

• Exporting pytorch model to onnx increases the model size