Documentation
Search…
Hugging Face
A Weights & Biases integration for Hugging Face's Transformers library: solving NLP, one logged run at a time!
The Hugging Face Transformers library makes state-of-the-art NLP models like BERT and training techniques like mixed precision and gradient checkpointing easy to use. The W&B integration adds rich, flexible experiment tracking and model versioning to interactive centralized dashboards without compromising that ease of use.

🤗 Next-level Hugging Face logging in 2 lines

1
from transformers import TrainingArguments, Trainer
2
3
args = TrainingArguments(... , report_to="wandb")
4
trainer = Trainer(... , args=args)
Copied!
Explore your experiment results in the W&B interactive dashboard

This guide covers

If you'd rather dive straight into working code, check out this Google Colab.

Getting started: track experiments

1) Install the wandb library and log in

Notebook
Command Line
1
!pip install wandb
2
3
import wandb
4
wandb.login()
Copied!
1
pip install wandb
2
wandb login
Copied!

2) Name the project

A Project is where all of the charts, data, and models logged from related runs are stored. Naming your project helps you organize your work and keep all the information about a single project in one place.
To add a run to a project simply set the WANDB_PROJECT environment variable to the name of your project. The WandbCallback will pick up this project name environment variable and use it when setting up your run.
Notebook
Command Line
1
%env WANDB_PROJECT=amazon_sentiment_analysis
Copied!
1
WANDB_PROJECT=amazon_sentiment_analysis
Copied!
Make sure you set the project name before you initialize the Trainer.
If a project name is not specified the project name defaults to "huggingface".

3) Log your training runs to W&B

This is the most important step: when defining your Trainer training arguments, either inside your code or from the command line, set report_to to "wandb" in order enable logging with Weights & Biases.
You can also give a name to the training run using the run_name argument.
Using TensorFlow? Just swap the PyTorch Trainer for the TensorFlow TFTrainer.
That's it! Now your models will log losses, evaluation metrics, model topology, and gradients to Weights & Biases while they train.
Notebook
Command Line
1
from transformers import TrainingArguments, Trainer
2
3
args = TrainingArguments(
4
# other args and kwargs here
5
report_to="wandb", # enable logging to W&B
6
run_name="bert-base-high-lr" # name of the W&B run (optional)
7
)
8
9
trainer = Trainer(
10
# other args and kwargs here
11
args=args, # your training args
12
)
13
14
trainer.train() # start training and logging to W&B
Copied!
1
python run_glue.py \ # run your Python script
2
--report_to wandb \ # enable logging to W&B
3
--run_name bert-base-high-lr \ # name of the W&B run (optional)
4
# other command line arguments here
Copied!

(Notebook only) Finish your W&B Run

If your training is encapsulated in a Python script, the W&B run will end when your script finishes.
If you are using a Jupyter or Google Colab notebook, you'll need to tell us when you're done with training by calling wandb.finish().
Notebook
1
trainer.train() # start training and logging to W&B
2
3
# post-training analysis, testing, other logged code
4
5
wandb.finish()
Copied!

4) Visualize your results

Once you have logged your training results you can explore your results dynamically in the W&B Dashboard. It's easy to compare across dozens of runs at once, zoom in on interesting findings, and coax insights out of complex data with flexible, interactive visualizations.

Advanced features

Turn on model versioning

Using Weights & Biases' Artifacts, you can store up to 100GB of models and datasets. Logging your Hugging Face model to W&B Artifacts can be done by setting a W&B environment variable called WANDB_LOG_MODEL to true.
Notebook
Command Line
1
%env WANDB_LOG_MODEL=true
Copied!
1
WANDB_LOG_MODEL=true
Copied!
Your model will be saved to W&B Artifacts as run-{run_name}.
Any Trainer you initialize from now on will upload models to your W&B project. Your model file will be viewable through the W&B Artifacts UI. See the Weights & Biases' Artifacts guide for more about how to use Artifacts for model and dataset versioning.

How do I save the best model?

If load_best_model_at_end=True is passed to Trainer, then W&B will save the best performing model checkpoint to Artifacts instead of the final checkpoint.

Loading a saved model

If you saved your model to W&B Artifacts with WANDB_LOG_MODEL, you can download your model weights for additional training or to run inference. You just load them back into the same Hugging Face architecture that you used before.
1
# Create a new run
2
with wandb.init(project="amazon_sentiment_analysis") as run:
3
4
# Connect an Artifact to the run
5
my_model_name = "run-bert-base-high-lr:latest"
6
my_model_artifact = run.use_artifact(my_model_name)
7
8
# Download model weights to a folder and return the path
9
model_dir = my_model_artifact.download()
10
11
# Load your Hugging Face model from that folder
12
# using the same model class
13
model = AutoModelForSequenceClassification.from_pretrained(
14
model_dir, num_labels=num_labels)
15
16
# Do additional training, or run inference
Copied!

Additional W&B settings

Further configuration of what is logged with Trainer is possible by setting environment variables. A full list of W&B environment variables can be found here.
Environment Variable
Usage
WANDB_PROJECT
Give your project a name
WANDB_LOG_MODEL
Log the model as artifact at the end of training (false by default)
WANDB_WATCH
Set whether you'd like to log your models gradients, parameters or neither
    gradients: Log histograms of the gradients (default)
    all: Log histograms of gradients and parameters
    false: No gradient or parameter logging
WANDB_DISABLED
Set to true to disable logging entirely (false by default)
WANDB_SILENT
Set to true to silence the output printed by wandb (false by default)
Notebook
Command Line
1
%env WANDB_WATCH=all
2
%env WANDB_SILENT=true
Copied!
1
WANDB_WATCH=all
2
WANDB_SILENT=true
Copied!

Customize wandb.init

The WandbCallback that Trainer uses will call wandb.init under the hood when Trainer is initialized. You can alternatively set up your runs manually by calling wandb.init before theTrainer is initialized. This gives you full control over your W&B run configuration.
An example of what you might want to pass to init is below. For more details on how to use wandb.init, check out the reference documentation.
1
wandb.init(project="amazon_sentiment_analysis",
2
name="bert-base-high-lr",
3
tags=["baseline", "high-lr"],
4
group="bert")
Copied!

Custom logging

Logging to Weights & Biases via the Transformers Trainer is taken care of by the WandbCallback (reference documentation) in the Transformers library. If you need to customize your Hugging Face logging you can modify this callback.

Issues, questions, feature requests

For any issues, questions, or feature requests for the Hugging Face W&B integration, feel free to post in this thread on the Hugging Face forums or open an issue on the Hugging Face Transformers GitHub repo.
Last modified 30d ago