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.
from transformers import TrainingArguments, Trainerargs = TrainingArguments(... , report_to="wandb")trainer = Trainer(... , args=args)
how to get started using W&B with Hugging Face Transformers to track your NLP experiments and
how to use advanced features of the W&B Hugging Face integration to get the most out of experiment tracking.
!pip install wandbimport wandbwandb.login()
pip install wandbwandb login
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.
If a project name is not specified the project name defaults to "huggingface".
This is the most important step: when defining your
Trainer training arguments, either inside your code or from the command line, set
"wandb" in order enable logging with Weights & Biases.
You can also give a name to the training run using the
That's it! Now your models will log losses, evaluation metrics, model topology, and gradients to Weights & Biases while they train.
from transformers import TrainingArguments, Trainerargs = TrainingArguments(# other args and kwargs herereport_to="wandb", # enable logging to W&Brun_name="bert-base-high-lr" # name of the W&B run (optional))trainer = Trainer(# other args and kwargs hereargs=args, # your training args)trainer.train() # start training and logging to W&B
python run_glue.py \ # run your Python script--report_to wandb \ # enable logging to W&B--run_name bert-base-high-lr \ # name of the W&B run (optional)# other command line arguments here
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
trainer.train() # start training and logging to W&B# post-training analysis, testing, other logged codewandb.finish()
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.
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
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.
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.
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.
# Create a new runwith wandb.init(project="amazon_sentiment_analysis") as run:# Connect an Artifact to the runmy_model_name = "run-bert-base-high-lr:latest"my_model_artifact = run.use_artifact(my_model_name)# Download model weights to a folder and return the pathmodel_dir = my_model_artifact.download()# Load your Hugging Face model from that folder# using the same model classmodel = AutoModelForSequenceClassification.from_pretrained(model_dir, num_labels=num_labels)# Do additional training, or run inference
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.
Give your project a name
Log the model as artifact at the end of training (
Set whether you'd like to log your models gradients, parameters or neither
%env WANDB_WATCH=all%env WANDB_SILENT=true
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 the
Trainer 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.
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.
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.