大于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模型的方法:
1. 把卷积和矩阵乘的权重(参数量大于某个阈值)替换为ConstantOfShape,从而显著缩小模型大小。
2. 利用onnxsim特性避免折叠(参数量大于某个阈值)ConstantOfShape算子。
需要onnxsim>=0.4.24
避免常量折叠产生大tensor的tile和ConstantOfShape算子的onnxsim命令:
onnxsim --no-large-tensor size_th in_model.onnx out_model.onnx
size_th类似为1KB, 1MB等
3. 对压缩后的模型进行优化和常量折叠后的模型删除ConstantOfShape算子,并替换为原来的权重。
该方法也可以用于其他优化,例如onnx infer shape和opset转换。因为大于2GB模型需要写回到文件才能调用官方infershape,而大于2GB模型opset转换无官方方法。
注意每个ConstantOfShape的value最好不一样,否则onnxsim会合并value和shape相同的ConstantOfShape。
项目代码:
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博客