ONNX transformation to cast int64 constants to int32 when possible

[fxmarty]

As per title.

Partially fixes #627 , we need to integrate this in this CLI and document + test.

Try with:

import onnx
from pathlib import Path
from optimum.onnx import model_to_int32

path = "/path/to/decoder_model.onnx"
model = onnx.load(path)

model = model_to_int32(model)

onnx.save(
    model,
    path,
    save_as_external_data=True,
    all_tensors_to_one_file=True,
    location=Path(path).name + "_data",
)

onnx.checker.check_model(path)

Inspect the original and transformed models "Slice" nodes.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[fxmarty]

Should test with ONNX Runtime, it's possible there's an issue (could avoid converting int64 to int32, and just clamp to avoid it):

(hf-inf) fxmarty@huggingface:~/hf_internship/optimum/gptj_tiny_bool_int32$ onnxsim decoder_model_int32.onnx decoder_model_int32_sim.onnx --input-shape input_ids:1,12 attention_mask:1,12
Your model contains "Tile" ops or/and "ConstantOfShape" ops. Folding these ops can make the simplified model much larger. If it is not expected, please 
specify "--no-large-tensor" (which will lose some optimization chances)
Simplifying...
Traceback (most recent call last):
  File "/home/fxmarty/anaconda3/envs/hf-inf/bin/onnxsim", line 8, in <module>
    sys.exit(main())
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/onnxsim/onnx_simplifier.py", line 434, in main
    model_opt, check_ok = simplify(
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/onnxsim/onnx_simplifier.py", line 186, in simplify
    model_opt_bytes = C.simplify(
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/onnxsim/onnx_simplifier.py", line 239, in Run
    sess = rt.InferenceSession(
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 347, in __init__
    self._create_inference_session(providers, provider_options, disabled_optimizers)
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 386, in _create_inference_session
    sess = C.InferenceSession(session_options, self._model_bytes, False, self._read_config_from_model)
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Type Error: Type parameter (Tind) of Optype (Slice) bound to different types (tensor(int32) and tensor(int64) in node (/transformer/h.0/attn/Slice_4).

[fxmarty]

Should be all good now, feel free to review.

[michaelbenayoun]

LGTM

[michaelbenayoun]

Does this replace aten ops related to Gather to a working version?
If so I would make it more explicit in the name.

[fxmarty]

No idea, I did not code it. I'm not sure in which case it is an useful transform.