How to export 🤗 Transformers Models to ONNX ?
ONNX is open format for machine learning models. It allows you to save your neural network's computation graph in a framework agnostic way, which might be particularly helpful when deploying deep learning models.
Indeed, businesses might have other requirements (languages, hardware, ...) for which the training framework might not be the best suited in inference scenarios. In that context, having a representation of the actual computation graph that can be shared accross various business units and logics across an organization might be a desirable component.
Along with the serialization format, ONNX also provides a runtime library which allows efficient and hardware specific execution of the ONNX graph. This is done through the onnxruntime project and already includes collaborations with many hardware vendors to seamlessly deploy models on various platforms.
Through this notebook we'll walk you through the process to convert a PyTorch or TensorFlow transformers model to the ONNX and leverage onnxruntime to run inference tasks on models from 🤗 transformers
Exporting 🤗 transformers model to ONNX
Exporting models (either PyTorch or TensorFlow) is easily achieved through the conversion tool provided as part of 🤗 transformers repository.
Under the hood the process is sensibly the following:
- Allocate the model from transformers (PyTorch or TensorFlow)
- Forward dummy inputs through the model this way ONNX can record the set of operations executed
- Optionally define dynamic axes on input and output tensors
- Save the graph along with the network parameters
How to leverage runtime for inference over an ONNX graph
As mentionned in the introduction, ONNX is a serialization format and many side projects can load the saved graph and run the actual computations from it. Here, we'll focus on the official onnxruntime. The runtime is implemented in C++ for performance reasons and provides API/Bindings for C++, C, C#, Java and Python.
In the case of this notebook, we will use the Python API to highlight how to load a serialized ONNX graph and run inference workload on various backends through onnxruntime.
onnxruntime is available on pypi:
- onnxruntime: ONNX + MLAS (Microsoft Linear Algebra Subprograms)
- onnxruntime-gpu: ONNX + MLAS + CUDA
Preparing for an Inference Session
Inference is done using a specific backend definition which turns on hardware specific optimizations of the graph.
Optimizations are basically of three kinds:
- Constant Folding: Convert static variables to constants in the graph
- Deadcode Elimination: Remove nodes never accessed in the graph
- Operator Fusing: Merge multiple instruction into one (Linear -> ReLU can be fused to be LinearReLU)
ONNX Runtime automatically applies most optimizations by setting specific SessionOptions.
Note:Some of the latest optimizations that are not yet integrated into ONNX Runtime are available in optimization script that tunes models for the best performance.
Forwarding through our optimized ONNX model running on CPU
When the model is loaded for inference over a specific provider, for instance CPUExecutionProvider as above, an optimized graph can be saved. This graph will might include various optimizations, and you might be able to see some higher-level operations in the graph (through Netron for instance) such as:
- EmbedLayerNormalization
- Attention
- FastGeLU
These operations are an example of the kind of optimization onnxruntime is doing, for instance here gathering multiple operations into bigger one (Operator Fusing).
Benchmarking PyTorch model
Note: PyTorch model benchmark is run on CPU
Benchmarking PyTorch & ONNX on CPU
Disclamer: results may vary from the actual hardware used to run the model
Quantization support from transformers
Quantization enables the use of integers (instead of floatting point) arithmetic to run neural networks models faster. From a high-level point of view, quantization works as mapping the float32 ranges of values as int8 with the less loss in the performances of the model.
Hugging Face provides a conversion tool as part of the transformers repository to easily export quantized models to ONNX Runtime. For more information, please refer to the following:
- Hugging Face Documentation on ONNX Runtime quantization supports
- Intel's Explanation of Quantization
With this method, the accuracy of the model remains at the same level than the full-precision model. If you want to see benchmarks on model performances, we recommand reading the ONNX Runtime notebook on the subject.