Package: fastt5
Reduce T5 model size by 3X and increase the inference speed up to 5X.InstallUsageDetailsFunctionalitiesBenchmarksOnnx modelQuantized onnx modelQuantized model scoresfurther improvementsLicenseGet HelpAcknowledgementsT5 models can be used for several NLP tasks such as summarization, QA, QG, translation, text generation, and more. Sequential text generation is naturally slow, and for larger T5 models it gets even slower.fastT5makes the T5 models inference faster by running it on onnxruntime. and it also decreases the model size by quantizing it.fastT5 library allows you to convert a pretrained T5 model to onnx, quantizes it, and gives the model as output which is running on an onnxruntime in a single line of code. You can also customize this whole process.InstallYou can install fastT5 from PyPI:pipinstallfastt5If you want to build from source:gitclonehttps://github.com/Ki6an/fastT5cdfastT5pip3install-e.UsageTheexport_and_get_onnx_model()method exports the given pretrained T5 model to onnx, quantizes it and runs it on the onnxruntime with default settings. The returned model from this method supports thegenerate()method of huggingface.If you don't wish to quantize the model then usequantized=Falsein the method.fromfastT5importexport_and_get_onnx_modelfromtransformersimportAutoTokenizermodel_name='t5-small'model=export_and_get_onnx_model(model_name)tokenizer=AutoTokenizer.from_pretrained(model_name)t_input="translate English to French: The universe is a dark forest."token=tokenizer(t_input,return_tensors='pt')tokens=model.generate(input_ids=token['input_ids'],attention_mask=token['attention_mask'],num_beams=2)output=tokenizer.decode(tokens.squeeze(),skip_special_tokens=True)print(output)to run the already exported model useget_onnx_model()you can customize the whole pipeline as shown in the below code example:fromfastT5import(OnnxT5,get_onnx_runtime_sessions,generate_onnx_representation,quantize)fromtransformersimportAutoTokenizermodel_or_model_path='t5-small'# Step 1. convert huggingfaces t5 model to onnxonnx_model_paths=generate_onnx_representation(model_or_model_path)# Step 2. (recommended) quantize the converted model for fast inference and to reduce model size.quant_model_paths=quantize(onnx_model_paths)# step 3. setup onnx runtimemodel_sessions=get_onnx_runtime_sessions(quant_model_paths)# step 4. get the onnx modelmodel=OnnxT5(model_or_model_path,model_sessions)...custom output pathsBy default, fastT5 creates amodelsfolder in the current directory and stores all the models. You can provide a custom path for a folder to store the exported models. And to run alreadyexported modelsthat are stored in a custom folder path: useget_onnx_model(onnx_models_path="/path/to/custom/folder/")fromfastT5importexport_and_get_onnx_model,get_onnx_modelmodel_name="t5-small"custom_output_path="/path/to/custom/folder/"# 1. stores models to custom_output_pathmodel=export_and_get_onnx_model(model_name,custom_output_path)# 2. run already exported models that are stored in custom path# model = get_onnx_model(model_name, custom_output_path)DetailsT5 is aseq2seqmodel (Encoder-Decoder), as it uses decoder repeatedly for inference, we can't directly export the whole model to onnx. We need to export the encoder and decoder separately.past_key_valuescontain pre-computed hidden-states (key and values in the self-attention blocks and cross-attention blocks) that can be used to speed up sequential decoding.models can only be exported with a constant number of inputs. Contrary to this, the decoder of the first step does not takepast_key_valuesand the rest of the steps decoders do. To get around this issue, we can create two decoders: one for the first step that does not takepast_key_valuesand another for the rest of the steps that utilize thepast_key_values.Next, we'll export all three models (encoder, decoder, init_decoder). And then quantize them, quantizing32bitto8bitshould give the 4x memory reduction. Since there is an extra decoder the model size reduces by 3x.Finally, we'll run the quantized model on onnx runtime.The inference is simple as the model supports thegenerate()method of huggingface.FunctionalitiesExport any pretrained T5 model to ONNX easily (withpast_key_values).The exported model supports beam search and greedy search and more viagenerate()method.Reduce the model size by3Xusing quantization.Up to5Xspeedup compared to PyTorch execution for greedy search and3-4Xfor beam search.BenchmarksThe benchmarks are the result of the T5-base model tested on English to French translation.Onnx modelThe following graph shows the latency of the quantized onnx model vs the PyTorch model for beam numbers varying from 1 to 9. The latencies shown here are for the mean of sequence lengths up to 130.The following heat map shows the X times faster which the ratio of latency of PyTorch to onnx model.
The onnx model outperforms most cases. however, the speed of the model drops for a longer sequence length.Quantized onnx modelQuantized models are lightweight models as mentioned earlier, these models have almost the same accuracy as the original model (quantized model scores are mentioned in the next section). Quantized onnx models have the lowest latency compared to both Onnx & PyTorch models.The model outperforms the PyTorch model by 5.7X for greedy search on average and 3-4X for beam search.Note : The results were generated onAMD EPYC 7B12, these results may vary from device to device. The Onnx models usually perform well on high-end CPUs with more cores.Quantized model scoresThe results were tested for English to French translation with beam search number of 3.Bleu_4METEORROUGE_Lt5-small (quant)0.2407690.2823420.468817t5-small (pytorch)0.2546010.2951720.492749t5-base (quant)0.2676060.3060190.499188t5-base (pytorch)0.2683460.3049690.503306t5-large (quant)0.2867260.3168450.503585t5-large (pytorch)0.2940150.3157740.508677Private HuggingFace Model Hub ModelsTheHuggingFace model hubsupports private models. To use a private, pre-trained version of T5 with fastT5 you first must have authenticated into HuggingFace ecosystem with$ transformers-cli login. Then, when using fastT5, there is an extra import and call:fromfastT5import(OnnxT5,get_onnx_runtime_sessions,generate_onnx_representation,quantize,set_auth_token)fromtransformersimportAutoTokenizerset_auth_token(True)# the rest of the code is the same as using a public modelIf you are unable to call$ transformers-cli loginor prefer to use your API Key, found athttps://huggingface.co/settings/token(orhttps://huggingface.co/organizations/ORG_NAME/settings/tokenfor organizations), you can pass that as a string toset_auth_token. Avoid hard-coding your API key into code by setting the environment variableHF_API_KEY=<redacted>, and then in code:importosfromfastT5import(OnnxT5,get_onnx_runtime_sessions,generate_onnx_representation,quantize,set_auth_token)fromtransformersimportAutoTokenizerauth_token=os.environ.get("HF_API_KEY")set_auth_token(auth_token)# code proceeds as normalfurther improvementscurrently the fastT5 library supports only the cpu version of onnxruntime, gpu implementation still needs to be done.graph optimization of the onnx model will further reduce the latency.Get HelpContact me atkiranr8k@gmail.comIf appropriate,open an issueon GitHubAcknowledgementsoriginal T5 papertransformersby huggingfaceonnxonnxruntimeby microsoftonnxt5@article{2019t5,author={ColinRaffelandNoamShazeerandAdamRobertsandKatherineLeeandSharanNarangandMichaelMatenaandYanqiZhouandWeiLiandPeterJ.Liu},title={ExploringtheLimitsofTransferLearningwithaUnifiedText-to-TextTransformer},journal={arXive-prints},year={2019},archivePrefix={arXiv},eprint={1910.10683},}