大于2GB模型onnxsim优化很耗时,容易挂掉,而且需要特别大的系统内存。
这里提出一种比较简单的优化大于2GB ONNX模型的方法:
基本思路是
把卷积和矩阵乘的权重(参数量大于某个阈值)替换为ConstantOfShape,从而显著缩小模型大小。
利用onnxsim特性避免折叠(参数量大于某个阈值)ConstantOfShape算子。
对压缩后的模型进行优化和常量折叠后,再删除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等
ConstantOfShape算子
创建一个ConstantOfShape算子
const_shape = (512, 512)
const_shape_init = onnx.helper.make_tensor(name="const_shape", data_type=onnx.TensorProto.INT64, dims=[
2], vals=const_shape, raw=False)
tensor_value_attr = onnx.helper.make_tensor("value", onnx.TensorProto.FLOAT, dims=[1], vals=[1])
constantofshape_node = onnx.helper.make_node(op_type="ConstantOfShape",
inputs=["const_shape"],
outputs=["input2"],
value=tensor_value_attr)