Optimize 🤗 Hugging Face models with Weights & Biases

Hugging Face provides tools to quickly train neural networks for NLP (Natural Language Processing) on any task (classification, translation, question answering, etc) and any dataset with PyTorch and TensorFlow 2.0.

Coupled with Weights & Biases integration, you can quickly train and monitor models for full traceability and reproducibility without any extra line of code! You just need to install the library, sign in, and your experiments will automatically be logged:

pip install wandb
wandb login

Note: To enable logging to W&B, set report_to to wandb in your TrainingArguments or script.

W&B integration with 🤗 Hugging Face can automatically:

• log your configuration parameters
• log your losses and metrics
• log gradients and parameter distributions
• log your model
• keep track of your code
• log your system metrics (GPU, CPU, memory, temperature, etc)

Here's what the W&B interactive dashboard will look like:

🛠️ Installation and set-up

We need the following 🤗 Hugging Face libraries:

• transformers contains an API for training models and many pre-trained models
• tokenizers is automatically installed by transformers and "tokenize" our data (ie it converts text to sequence of numbers)
• datasets contains a rich source of data and common metrics, perfect for prototyping

We also install wandb to automatically instrument our training.

We finally make sure we're logged into W&B so that our experiments can be associated to our account.

💡 Configuration tips

W&B integration with Hugging Face can be configured to add extra functionalities:

• auto-logging of models as artifacts: just set environment varilable WANDB_LOG_MODEL to true
• log histograms of gradients and parameters: by default gradients are logged, you can also log parameters by setting environment variable WANDB_WATCH to all
• set custom run names with run_name arg present in scripts or as part of TrainingArguments
• organize runs by project with the WANDB_PROJECT environment variable

For more details refer to W&B + HF integration documentation.

Let's log every trained model.

🚅 Training a new model the quick way!

When working on a new problem, you should always check the summary of task as there will often be a script that can already solve your task. At a minimum they will be a great source of inspiration for your own custom pipeline.

Let's use the Hugging Face script responsible for training on any GLUE task, such as sequence classification.

These scripts are automatically instrumented with logging when wandb is installed and logged in.

Just set report_to to wandb to enable logging through W&B.

Note: This cell can take up to 5 minutes to run.

You just trained a model and can now visualize your metrics in your dashboard!

Just click the run page link that appears in the wandb section output of the cell above, just before training launches and just after ir finishes.

It should look something like this:

In addition, your model files have been saved and versioned, along with associated metadata (evaluation & training metrics).

Just check the "Artifacts" tab on your run page -- it's the one with the "stacked pucks" icon.

🔬 Advanced usage & custom training

Let's create our own logic for a more customized training.

✏️ Preparing a dataset

The dataset will vary based on the task you work on. Let's work on sequence classification!

Our dataset will be composed of sentences and their associated classes. For example if you wanted to identify the subject of a conversation, you could create a dataset such as:

The objective of our trained model will be to correctly identify the class associated to new sentences.

🔎 Finding a dataset

If you don't have the right dataset, you can always explore the Datasets Hub. The "topic classification" category contains many datasets suitable for prototyping this model.

We select "Yahoo! Answers Topic Classification" and visualize it with the Datasets viewer.

Each topic number reprersent a unique subject:

• 0:"Society & Culture"
• 1:"Science & Mathematics"
• 2:"Health"
• etc…

They correspond to the output we will try to predict from the model.

🏷️ Loading a dataset

Any dataset from the Datasets Hub can easily be loaded and is automatically downloaded if not present locally.

Just printing our dataset object gives us a lot of information.

We can easily access any element.

str2int and int2str help us go from class label to their integer mapping.

For our topic classification task, we use question_title as input and try to predict topic.

This particular dataset is split between 10 different topics, that will be represented by 10 classes from our model output.

The "topic" class needs to be renamed to "labels" for the Trainer to find it.

⚙️ Tokenizing the dataset

In order to train a neural network, we need to convert our inputs to numbers:

• the tokenizer divides a sequence of characters into tokens, ie sub-sequences (such as words, characters, sub-words…)
• each unique token is mapped to a unique integer

There are many types of tokenizers. 🤗 Transformers can auto-select the right Tokenizer associated to a specific model.

The tokenizer let us quickly preprocess our data.

The tokenizer can quickly process an entire dataset and cache the results locally to avoid any future tokenization of the same data.

We leverage dataset.map(fn) function which can efficiently apply any function to a dataset. We also take advantage of batch processing which is supported by the tokenizer and makes the operation even faster.

We truncate the data to the max length supported by the model. During training, we will pass the tokenizer to pad inputs to the longest sequence of the batch (the model requires same length inputs in a single batch).

Our dataset now contains new keys: input_ids (tokens) and attention_mask (needed for certain models).

✨ Loading a model

Plenty of models are available and can be explored on the Model Hub.

Once a model has been selected, it can be automatically loaded and adapted to one of its supported tasks.

In this case, we are loading a pre-trained network to which a custom head has been added for sequence classification and presents 10 classes corresponding to the possible topics of this dataset.

Let's make a function to return the topic prediction from a sample question.

Let's test a prediction on a sample sentence.

Obviously the model has not been trained yet so the results are still random.

🎉 Training the model

We now need to fine-tune the model based on our dataset.

The Trainer class let us easily train a model and is very flexible.

Note: set report_to to wandb in TrainingArguments to enable logging through W&B.

For more customization, refer to TrainingArguments documentation.

We can optionally define metrics to calculate in addition to the loss through the compute_metrics function.

Several metrics are readily available from the datasets library to monitor model performance.

The Trainer handles all the training & evaluation logic.

We can verify that we initially have an accuracy of about 10% (random predictions over 10 classes).

We start training by simply calling train().

We can monitor losses, metrics, gradients and parameters as the model trains.

When training is complete, our model is logged and versioned along with its performance as metadata.

We can now use the trained model for better predictions.

When we want to close our W&B run, we can call wandb.finish() (mainly useful in notebooks, called automatically in scripts).

Once you're happy with a model, don't forget to share it with the word on the Model Hub!

📚 Resources

• Hugging Face and W&B integration documentation contains a few tips for taking most advantage of W&B
• 🤗 Transformers documentation is extremely thorough and full of examples

❓ Questions about W&B

If you have any questions about using W&B to track your model performance and predictions, please reach out to the slack community.