Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Path: blob/main/examples/summarization_ort.ipynb
Views: 2535
How to fine-tune a T5 model with ONNX Runtime
This notebook is largely inspired by the summarization notebook of Transformers which takes PyTorch as backend for fine tuning.
Here you will use the ORTSeq2SeqTrainer
class in Optimum library and take ONNX Runtime as backend to accelerate the training.
In this notebook, we will walk through the fine-tuning of T5-small model in the 🤗 Transformers for a summarization task. We will use the XSum dataset (for extreme summarization) which contains BBC articles accompanied with single-sentence summaries, and the training as well as inference will be done by leveraging ORTSeq2SeqTrainer
in Optimum!
Let's speed the training up!
Dependencies
To use ONNX Runtime for training, you need a machine with at least one NVIDIA GPU.
ONNX Runtime training module need to be properly installed before launching the notebook! Please follow the instruction in Optimum's documentation to set up your environment.
Check your GPU:
If you're opening this Notebook on colab, you will probably need to install 🤗 Optimum, 🤗 Transformers, 🤗 Datasets and 🤗 evaluate. Uncomment the following cell and run it.
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
[Optional] If you want to share your model with the community and generate an inference API, there are a few more steps to follow.
First you have to store your authentication token from the Hugging Face website (sign up here if you haven't already!) then execute the following cell and input your username and password:
Then you need to install Git-LFS. Uncomment the following instructions:
Make sure your version of Transformers is at least 4.15.0:
Setup
We also quickly upload some telemetry - this tells us which examples and software versions are getting used so we know where to prioritize our maintenance efforts. We don't collect (or care about) any personally identifiable information, but if you'd prefer not to be counted, feel free to skip this step or delete this cell entirely.
Loading the dataset
We will use the 🤗 Datasets library to download the data and get the metric we need to use for evaluation (to compare our model to the benchmark). This can be easily done with the functions load_dataset
and load_metric
.
[Optional] To get a sense of what the data looks like, the following function will show some examples picked randomly in the dataset.
The metric is an instance of datasets.Metric
:
Preprocessing the data
Before we can feed those texts to our model, we need to preprocess them. This is done by a 🤗 Transformers Tokenizer
which will (as the name indicates) tokenize the inputs (including converting the tokens to their corresponding IDs in the pretrained vocabulary) and put it in a format the model expects, as well as generate the other inputs that the model requires.
To do all of this, we instantiate our tokenizer with the AutoTokenizer.from_pretrained
method, which will ensure:
we get a tokenizer that corresponds to the model architecture we want to use,
we download the vocabulary used when pretraining this specific checkpoint.
That vocabulary will be cached, so it's not downloaded again the next time we run the cell.
To prepare the targets for our model, we need to tokenize them inside the as_target_tokenizer context manager. This will make sure the tokenizer uses the special tokens corresponding to the targets:
If you are using one of the five T5 checkpoints we have to prefix the inputs with "summarize:" (the model can also translate and it needs the prefix to know which task it has to perform).
We can then write the function that will preprocess our samples. We just feed them to the tokenizer
with the argument truncation=True
. This will ensure that an input longer that what the model selected can handle will be truncated to the maximum length accepted by the model. The padding will be dealt with later on (in a data collator) so we pad examples to the longest length in the batch and not the whole dataset.
To apply this function on all the pairs of sentences in our dataset, we just use the map
method of our dataset
object we created earlier. This will apply the function on all the elements of all the splits in dataset
, so our training, validation and testing data will be preprocessed in one single command.
Even better, the results are automatically cached by the 🤗 Datasets library to avoid spending time on this step the next time you run your notebook. The 🤗 Datasets library is normally smart enough to detect when the function you pass to map has changed (and thus requires to not use the cache data). For instance, it will properly detect if you change the task in the first cell and rerun the notebook. 🤗 Datasets warns you when it uses cached files, you can pass load_from_cache_file=False
in the call to map
to not use the cached files and force the preprocessing to be applied again.
Note that we passed batched=True
to encode the texts by batches together. This is to leverage the full benefit of the fast tokenizer we loaded earlier, which will use multi-threading to treat the texts in a batch concurrently.
Fine-tuning the model
Now that our data is ready, we can download the pretrained model and fine-tune it. Since our task is of the sequence-to-sequence kind, we use the AutoModelForSeq2SeqLM
class to fist load the PyTorch model. Like with the tokenizer, the from_pretrained
method will download and cache the model for us.
Note that we don't get a warning like in our classification example. This means we used all the weights of the pretrained model and there is no randomly initialized head in this case.
To instantiate a ORTSeq2SeqTrainer
, we will need to define three more things. The most important is the ORTSeq2SeqTrainingArguments
, which is a class that contains all the attributes to customize the training. You can also use Seq2SeqTrainingArguments
in Transformers, but ORTSeq2SeqTrainingArguments
enables more optimized features of ONNX Runtime. It requires one folder name, which will be used to save the checkpoints of the model, and all other arguments are optional:
Here we set the evaluation to be done at the end of each epoch, tweak the learning rate, use the batch_size
defined at the top of the cell and customize the weight decay. Since the ORTSeq2SeqTrainer
will save the model regularly and our dataset is quite large, we tell it to make three saves maximum. Lastly, we use the predict_with_generate
option (to properly generate summaries) and activate mixed precision training (to go a bit faster).
The last argument to setup everything so we can push the model to the Hub regularly during training. Remove it if you didn't follow the installation steps at the top of the notebook. If you want to save your model locally in a name that is different than the name of the repository it will be pushed, or if you want to push your model under an organization and not your name space, use the hub_model_id argument to set the repo name (it needs to be the full name, including your namespace: for instance "optimum/t5-large-finetuned-xsum"
).
Then, we need a special kind of data collator, which will not only pad the inputs to the maximum length in the batch, but also the labels:
The last thing to define for our ORTSeq2SeqTrainer
is how to compute the metrics from the predictions. We need to define a function for this, which will just use the metric
we loaded earlier, and we have to do a bit of pre-processing to decode the predictions into texts:
Then we just need to pass all of this along with our datasets to the ORTSeq2SeqTrainer
:
We can now finetune our model by just calling the train
method:
You can now upload the result of the training to the Hub, just execute this instruction:
You will also be able to save your fine-tuned model as PyTorch or ONNX model in the output_dir
that you set in ORTSeq2SeqTrainer
:
Evaluating your model
Evaluate the performance of the model that you just fine-tuned with the validation dataset that you've passed to ORTSeq2SeqTrainer
by just calling the evaluate
method.
If you set inference_with_ort=True
, the inference will be done with ONNX Runtime backend. Otherwise, the inference will take PyTorch as backend.
Extended reading
Now check your trained ONNX model with Netron, and you might notice that the computation graph is under optimizatiom. Want to accelerate even more?
Check the graph optimizers and quantizers of Optimum🤗!