【昇腾】常用训练框架导出ONNX，（高能）大于2GB的模型如何导出ONNX

原创

已于 2025-01-07 12:18:52 修改 · 1.9k 阅读

34 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能

于 2025-01-06 20:48:09 首次发布

虽然大模型兴起后，各种大模型加速方案兴起，但还是有很多小规模（2B以下）的模型基于ONNX方案，特别是传统CV、语音、NLP等场景，也包括Transformer结构的小规模参数模型。基于PyTorch、MindSpore和TensorFlow等框架将模型训练完成后，将模型转成ONNX，部署在中心或者边缘AI产品上。下面就讲下各训练框架训练好的模型如何保存为ONNX文件。由于受protobuf限制，通常只能保存小于2GB的ONNX模型文件。如何解决这个问题呢，今天给出一个方法。

PyTorch导出ONNX

import torch
from torchvision.models import resnet50, ResNet50_Weights

# 构造输入及shape
dummy_input = torch.randn(1, 3, 224, 224)

weights = ResNet50_Weights.DEFAULT
model = resnet50(weights=weights)
model.eval()

input_names = [ "input_data" ] 
output_names = [ "output_data" ]

torch.onnx.export(
        model,                # pytorch网络模型
        dummy_input,          #