pytorch

4f4e2c1c - Add constant node sizes to proto size calculation (#111097)

Add constant node sizes to proto size calculation (#111097) Fixes #110982 https://github.com/pytorch/pytorch/pull/62257 deprecated `torch.onnx.export(use_external_data_format: bool=...)` argument, but it seems the introduced `EncoderBase::GetGraphProtoSize` has a bug and doesn't detect models > 2GB when onnx Constant nodes are large (and responsible for the size overflow) This PR adds the constant node to the total size of the model, along with initializers. In python, what we need to do is: ```python import onnx def compute_tensor_size(tensor): # Compute the size of the tensor based on its shape and data type size = tensor.size * tensor.itemsize return size def sum_constant_and_initializer_sizes(model_path): # Load the ONNX model model = onnx.load(model_path) total_size = 0 initializer_size = 0 constant_size = 0 # Compute the size of constant nodes for node in model.graph.node: if node.op_type == 'Constant': constant_value = node.attribute[0].t # Convert constant value to numpy array constant_array = onnx.numpy_helper.to_array(constant_value) # Compute the size of the constant tensor tensor_size = compute_tensor_size(constant_array) total_size += tensor_size constant_size += tensor_size # Compute the size of initializer nodes that are not graph inputs for initializer in model.graph.initializer: if initializer.name not in [input.name for input in model.graph.input]: # Convert the shape and data type information to calculate size # tensor = onnx.helper.tensor_value_info_to_tensor(input) tensor = onnx.numpy_helper.to_array(initializer) tensor_size = compute_tensor_size(tensor) total_size += tensor_size initializer_size += tensor_size return total_size, constant_size, initializer_size model_path = '/path/to/model.onnx' total_size, constant_size, initializer_size = sum_constant_and_initializer_sizes(model_path) print("Total size of constant nodes in bytes:", constant_size) print("Total size of initializer nodes (excluding graph inputs) in bytes:", initializer_size) print("Total size of constant and initializer nodes (excluding graph inputs) in bytes:", total_size) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/111097 Approved by: https://github.com/justinchuby, https://github.com/zhipenghan