In the follow up PR we will remove the checks altogether from a discussion with Aaron: we should not check it here.

Thanks GPT

So the trick is the 2GB limitation only applies when serializing, but not on in memory ModelProto.

Makes me think _export_onnx should return ModelProto instead of the serialized string. But that is not supported by pybind. We can probably create a c++ pybind api that returns ModelProto and initializers in separate serialized strings, then in python deserialize and combine them together. Only python api is exposed and used here which gives us ModelProto directly.

The benefit is we skip checking size, nor write to disk. What do you think?

Do we return a list of initializers?

Taking this further, we don’t even need to pass in initializers to _export_onnx. We can serialize them ourselves outside with onnx. This way we don’t need to change the PyTorch c++ implementation.

Tested with

def _add_initializers(model_proto: onnx.ModelProto, initializers: Mapping[str, torch.Tensor]):
    tensor_protos = []
    for name, tensor in initializers.items():
        print(name, "0")
        tensor_numpy = tensor.detach().numpy()
        print(name)
        tensor_proto = onnx.helper.make_tensor(
            name=name,
            data_type=onnx.helper.np_dtype_to_tensor_dtype(tensor_numpy.dtype),
            dims=tensor_numpy.shape,
            vals=tensor_numpy,
        )
        print(name, "done1")
        tensor_protos.append(tensor_proto)
    model_proto.graph.initializer.extend(tensor_protos)

But onnx.helper.make_tensor is very slow.

I think returning a list of TensorProtos can work. But for now it seems to me allowing the compatibility may be nice.

*compatibility with torch2.0

Because we will move away from torchscript eventually

Let's try make_tensor(...torch_tensor.data_ptr().to_bytes(torch_tensor.element_size() * torch_tensor.numel(), byteorder=sys.byteorder), raw=True) and if that doesn't work or is still slow then we go back to initial solution.

Should be

import ctypes
torch_tensor = torch.tensor([2, 3])
raw_data = bytes(ctypes.c_ubyte*torch_tensor.element_size()*torch_tensor.numel()).from_address(torch_tensor.data_ptr())
tensor_proto = make_tensor(..., vals=raw_data, raw=True)

I was misled by someone doing data_ptr().to_bytes(), but what that really does was converting the pointer integer itself to bytes... lol

I'm not sure if it is worth it, looks hacky, but this should resemble what _export_onnx is doing on c++ side. If that is still slow then there is nothing more we can do.

#1004

For shape inference, can we still load shape and element type from model file (not initializer files) and then run infer_shape?

We could, but we also don’t need to because PyTorch supplies all the shape info.

A drawback due to onnx/onnx#5487, we don't have much inner node shape info left now that modules are functions.

Humm, maybe increase to 1.8 GB? I never see a model > 100MB without initializers.

Done

Whether or not storing initializers should be controlled by a user flag. Assume that I export a 1GB model on remote machine. I want to visualize it locally. I really don't want to download its initializers with home internet. If this flag can be turned on, I will be able to just download the structure of model and debug faster.

I think once the user get the model proto, they can do whatever they want (aka remove all the data)? A user has full control when they get the dynamo export output as an object.

Further yet include_initializers is already an argument

Agree with @wschin 's goal and @justinchuby 's explanation. A thing to consider for ExportOutput.save or ExportOutputSerializer.