文章提出了针对大于2GB的ONNX模型的优化方法，包括将卷积和矩阵乘权重替换为ConstantOfShape，利用onnxsim避免大张量的优化，以及在优化后删除ConstantOfShape并恢复原始权重。这种方法可以减少内存需求和优化时间，适用于stablediffusionunet等大模型。需要注意ConstantOfShape的value应保持唯一，以及处理动态shape的问题。
大于2GB模型onnxsim优化很耗时，容易挂掉，而且需要特别大的系统内存。
此外直接优化大于2GB模型可能报错：
model with IR version >= 3 must be specify opset_import for ONNX。
onnx.onnx_cpp2py_export.checker.ValidationError: The model does not have an ir_version set properly.
这里提出一种比较简单的优化大于2GB ONNX模型的方法：
把initializer换成模型的输入，onnxsim后再把input换成原来的initializer。
该方法也可以用于其他优化，例如onnx infer shape和opset转换。因为大于2GB模型需要写回到文件才能调用官方infershape，而大于2GB模型opset转换无官方方法。
项目代码：
https://github.com/luchangli03/onnxsim_large_model
该方法可以成功用于stable diffusion unet, llama, chatglm等导出的onnx的优化。
该方法可以显著降低onnxsim大模型需要的内存，以及优化时间。
需要注意几点，stable diffusion unet可以采用上面的方法压缩，然后设置压缩模型的输入shape并进行onnxsim优化。但是可能一次优化并不能消除所有的shape算子等动态shape 并且可能出现time_step从[1]变成[-1]的bug。需要重新设置一次输入shape再onnxsim优化一次即可消除所有动态shape算子。最后再进行解压缩。
如果明明设置了静态shape，onnxsim还是优化不完全，可以先删除旧的shape value info，再调用onnx infer shape工具infer shape，再使用onnxsim, 参考：
onnx模型图优化/模型修改_nodes in a graph must be topologically sorted_Luchang-Li的博客-CSDN博客
扫一扫
3694
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
