How to edit onnx larger than 2GB with onnx_graphsurgeon?

[1049451037]

Description

I have a UNet larger than 2GB. When I use the TensorRT Diffusion demo to optimize the onnx, it produce error when running opt.infer_shapes.

ValueError: Message onnx.ModelProto exceeds maximum protobuf size of 2GB: 2602836132

Is there any solution for this case?

Relevant Files

https://github.com/NVIDIA/TensorRT/tree/release/8.5/demo/Diffusion

[rajeevsrao]

For ONNX models larger than 2GB, you could use the infer_shapes_path API to write the inferred model to file. Ex:

if onnx_graph.ByteSize() > 2147483648:
    onnx.shape_inference.infer_shapes_path(self.onnx_path, self.onnx_opt_path)
    onnx_graph = onnx.load(self.onnx_opt_path)
else:
    onnx_graph = onnx.shape_inference.infer_shapes(onnx_graph)

[1049451037]

Moreover, I have an onnx that is not able to inference shape by onnxruntime... (There is a plugin node in the onnx file...) Is there any workaround?

[ttyio]

closing since there is already solution to infer shapes with onnxmodels larger than 2Gb, for the remaining plugin node issue, could you ask in the onnxruntime, there are more onnxruntime experts, thanks!