microsoft / onnxscript

ONNX Script enables developers to naturally author ONNX functions and models using a subset of Python.
https://onnxscript.ai/
MIT License
280 stars 53 forks source link

`proto2python` should not codegen large tensor data as 'make_tensor' #806

Open BowenBao opened 1 year ago

BowenBao commented 1 year ago

Otherwise generated python code is too large with model parameters defined explicitly.

Example with gpt2 model.

import torch
import transformers
import onnxscript

model = transformers.GPT2Model.from_pretrained("gpt2")
tokenizer = transformers.GPT2Tokenizer.from_pretrained("gpt2")
model.eval()

inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")

onnx_model = torch.onnx.dynamo_export(model, **inputs).model_proto
import onnx
onnx.save(onnx_model, "gpt2.onnx")
code = onnxscript.proto2python(onnx_model)
with open("gpt2_model.py", "w") as f:
    f.write(code)

Produces

2.6G gpt2_model.py
487M gpt2.onnx
gramalingam commented 1 year ago

Hi Bowen, for now, you can do this by explicitly using the save_as_external_data option when you call onnx.save. This converts the model to use external data. Calling proto2python on this model will do the right thing. Eg., use a save call as below:

def save(model, folder):
    datafile = os.path.join(folder, "external_data.onnx")
    modelfile = os.path.join(folder, "model.onnx")
    onnx.save(model, modelfile, save_as_external_data=True, all_tensors_to_one_file=True, size_threshold=64, convert_attribute=True, location=datafile)
    pymodel = proto2python(model, use_operators=True, inline_const=True)
    with open(os.path.join(folder, "model.py"), "w") as f:
        f.write(pymodel)
gramalingam commented 1 year ago

The inline_const option in proto2python will try to produce more readable code by inlining small constants (typically shapes).