Add possibility to output attentions from the ONNX models

Feature request

Hello!

I am interested in the possibility of retrieving not only the logits from the ONNX exported model, but also the attentions (which can be set via output_attentions=True in transformers). I have tried setting output_attentions=True in the config.json file of the model repository, but an exception is raised when the onnxruntime InferenceSession is initialized.

Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from ort_model_optimized.onnx failed:/private/var/folders/24/8k48jl6d249_n_qfxwsl6xvm0000gn/T/pip-req-build-u762x_rp/onnxruntime/core/graph/graph.cc:1236 void onnxruntime::Graph::InitializeStateFromModelFileGraphProto() This is an invalid model. Graph output (attentions) does not exist in the graph

Motivation

I am using attention scores for calculating the importance, but it is currently impossible to extract them in our production environment as we are using ONNX-exported transformers models.

Your contribution

If you could guide me a bit on what to do, I could submit a PR.

[jegork]

Hello @JingyaHuang, after taking a deeper look, it turned out that attention fusion was causing the aforementioned problem, so just using OptimizationConfig(disable_attention=True) worked well enough (although the attention fusion is not working, the remaining fusion optimization operations still work correctly).

Thank you for your time and guidance!

[JingyaHuang]

Hi @jegork , glad to hear that you have found a way to solve it. By setting optimize_with_onnxruntime_only=True you will turn off the whole fusion optimization, have you tried precising FusionOptions which can be set within your OptimizationConfig to locate which fusion leading to the failure? This might help enabling some other non-fatal fusions.