Saving and Restoring Machine Learning Models with W&B
In this article, we'll look at how to save and restore your machine learning models with Weights & Biases.
Created on March 27|Last edited on November 14
Comment
Introduction
In this notebook, we'll look at how to save and restore your machine learning models with Weights & Biases.
W&B lets you save everything you need to reproduce your models - weights, architecture, predictions, and code to a safe place.
This is useful because you don’t have to re-train your models, you can simply view their performance days, weeks, or even months later. Before you're ready to deploy, you can compare the performance of all the models you trained in the previous months and restore the best-performing one.
If you'd like to follow along, I've created a Colab to provide working code samples and an easy demo.
Save your machine learning model
There are two ways to save a file to associate with a run.
- Use wandb.save(filename).
- Put a file in the wandb run directory, and it will get uploaded at the end of the run.
If you want to sync files as they're being written, you can specify a filename or glob in wandb.save.
# "model.h5" is saved in wandb.run.dir & will be uploaded at the end of trainingmodel.save(os.path.join(wandb.run.dir, "model.h5"))# Save a model file manually from the current directory:wandb.save('model.h5')# Save all files that currently exist containing the substring "ckpt":wandb.save('../logs/*ckpt*')# Save any files starting with "checkpoint" as they're written to:wandb.save(os.path.join(wandb.run.dir, "checkpoint*"))
You can view your saved models by navigating to a run page, clicking on the Files tab, then clicking on your model file. See an example here.
Restore your machine learning model
Now let's look at how to restore a file, such as a model checkpoint, into your local run folder to access in your script.
Common use cases:
- restore the model architecture or weights generated by past runs
- resume training from the last checkpoint in the case of failure (see the section on resuming for crucial details)
# restore the model file "model.h5" from a specific run by user "lavanyashukla"# in project "save_and_restore" from run "10pr4joa"best_model = wandb.restore('model.h5', run_path="lavanyashukla/save_and_restore/10pr4joa")# use the "name" attribute of the returned object if your framework expects a filename, e.g. as in Kerasmodel.load_weights(best_model.name)
Related Articles
How to save and load models in PyTorch
This article is a machine learning tutorial on how to save and load your models in PyTorch using Weights & Biases for version control.
Using Artifacts to Build an End to End ML Pipeline
Data collection to model deployment and back again!
An Introduction to Weights & Biases Artifacts
This article introduces Weights & Biases Artifacts, which offer lightweight dataset versioning, and gives an easy-to-follow example so that you can see how they work.
Organize Your Machine Learning Pipelines With Artifacts in Weights & Biases
In this article, we'll look at how to use W&B Artifacts to store and keep track of datasets, models, and evaluation results across machine learning pipelines.
Add a comment
Iterate on AI agents and models faster. Try Weights & Biases today.