Model Management

We are actively investing in new Model Management features within W&B. Please contact us with questions or suggestions at [email protected]
Prefer code? Jump right into the companion colab notebook now!
A Model Registry is a system of record for organizing ML Models - often serving as an interface between model producers and consumers. This guide will show you how to use W&B as a Model Registry to track and report on the complete workflow of developing a model:
1.Model Catalog & Versioning: Save and restore every version of your model & learned parameters - organize versions by use case and objective.
2.Model Metadata: Track training metrics, assign custom metadata, and document rich markdown descriptions of your models.
3.Model Lineage Tracking & Reproducibility: Track the exact code, hyperparameters, & training dataset used to produce the model.
4.Model Lifecycle: Promote promising models to positions like "staging" or "production" - allowing downstream users to fetch the best model automatically.
5.Model Reporting: Create a Report dashboard to summarize model progression and performance over time
Data Model & Terminology
First, let's define the data model & terminology used throughout this guide:
Model Version: a package of data & metadata describing a trained model.
Model Artifact: a sequence of logged Model Versions - often tracking the progress of training.
Model Portfolio: a selection of linked Model Versions - often representing all the candidate models for a single modeling use case or task
Hint: A Model Version will always belong to one and only one Model Artifact, yet may belong to zero or more Model Portfolios.
For those familiar with W&B Artifacts: a Model is exactly an Artifact with type="model", a Model Version is an Artifact Version belonging to such an Artifact, and a Model Portfolio is an Artifact Portfolio of type="model".
In W&B a Model Version is an immutable directory of data; it is up to you to decide what files & formats are appropriate to store (and restore) your model architecture & learned parameters. Typically you will want to store whatever files are produced from the serialization process provided by your modeling library (eg PyTorch & Keras).
Furthermore, a Model Artifact is a sequence of Model Versions. Model Artifact can alias specific versions so that downstream consumers can pin to such aliases. It is extremely common for a W&B Run to produce many versions of a model while training (periodically saving checkpoints). Using this approach, each individual model being trained by the Run corresponds to its own Model Artifact, and each checkpoint corresponds to its own Model Version of the respective Model Artifact. View an example Model Artifact!
Finally, a Model Portfolio is a set links to Model Versions. A Model Portfolio can be accessed exactly like act like Model Artifacts (identified by [[entityName/]/projectName]/portfolioName:alias), however it acts more like a folder of "bookmarks" - where each "version" of a Model Portfolio is actually a link to an Model Version belonging to a Model Artifact of the same type. A Model Version may be linked to any number of Model Portfolios. Typically you will create a Model Portfolio for each of your use cases / modeling tasks and use aliases like "production" or "staging" to denote versions with special purposes. View an example Model Portfolio!
While developing an ML Model, you will likely have dozens, hundreds, or even thousands of Runs which produce Model Versions - they may come from notebooks, remote training jobs, CI/CD pipelines, etc... Most likely, not all of those models are great; often you are iterating on scripts, parameters, architectures, preprocessing logic and more. The separation of Artifacts and Portfolios allows you to produce a massive number of Artifacts (think of them like "draft models"), and periodically link your high performing versions to a the curated Model Portfolio. Then, use aliases to mark which Version in a Portfolio is at which stage in the lifecycle. Each person in your team can collaborate on a single use case (Portfolio), while having the freedom to explore and experiment without polluting namespaces or conflicting with others' work.
Workflow
Please add artifact-portfolios to your bio in order to enable beta features required to complete this guide.
Now we will walk through a canonical workflow for producing, organizing, and consuming trained models:
1.​Create a new Model Portfolio​
2.​Train & log Model Versions​
3.​Link Model Versions to the Portfolio​
4.​Using a Model Version​
5.​Evaluate Model Performance​
6.​Promote a Version to Production​
7.​Use the Production Model for Inference​
8.​Build a Reporting Dashboard​
A companion colab notebook is provided which covers step 2-3 in the first code block and steps 4-6 in the second code block.
1. Create a new Model Portfolio
First, create a Model Portfolio to hold all the candidate models for your particular modeling task. In this tutorial, we will use the classic MNIST Dataset - 28x28 grayscale input images with output classes from 0-9. The video below demonstrates how to create a new Portfolio:
1.Visit your Project's Artifact Browser: wandb.ai/<entity>/<project>/artifacts
2.Click the + icon on the bottom of the Artifact Browser Sidebar
3.Select Type: model, Style: Portfolio, and enter a name. In our case MNIST Grayscale 28x28. Remember, a Portfolio should map to a modeling task - enter a unique name that describes the use case.
2. Train & log Model Versions
Next, you will log a model from your training script:
1.(Optional) Declare your dataset as a dependency so that it is tracked for reproducibility and audibility
2.Serialize your model to disk periodically (and/or at the end of training) using the serialization process provided by your modeling library (eg PyTorch & Keras).
3.Add your model files to an Artifact of type "model"
Note: We use the name f'mnist-nn-{wandb.run.id}'. While not required, it is advisable to name-space your "draft" Artifacts with the Run id in order to stay organized
4.Log your model
Note: If you are logging multiple versions, it is advisable to add an alias of "best" to your Model Version when it outperforms the prior versions. This will make it easy to find the model with peak performance - especially when the tail end of training may overfit!
5.(Optional) Log training metrics associated with the performance of your model during training. 
Note: The data logged immediately after logging your Model Version will automatically be associated with that version.
By default, you should use the native W&B Artifacts API to log your serialized model. However, since this pattern is so common, we have provided a single method which combines serialization, Artifact creation, and logging. See the "(Beta) Using log_model" tab for details.
Using Artifacts
(Beta) Using `log_model`
Declare Dataset Dependency
1
import wandb
2
​
3
# Always initialize a W&B run to start tracking
4
wandb.init()
5
​
6
# (Optional) Declare an upstream dataset dependency
7
# see the `Declare Dataset Dependency` tab for
8
# alternative examples.
9
dataset = wandb.use_artifact("mnist:latest")
10
​
11
# At the end of every epoch (or at the end of your script)...
12
# ... Serialize your model
13
model.save("path/to/model.pt")
14
# ... Create a Model Version
15
art = wandb.Artifact(f'mnist-nn-{wandb.run.id}', type="model")
16
# ... Add the serialized files
17
art.add_file("path/to/model.pt", "model.pt")
18
# ... Log the Version
19
if model_is_best:
20
    # If the model is the best model so far, add "best" to the aliases
21
    wandb.log_artifact(art, aliases=["latest", "best"])
22
else:
23
    wandb.log_artifact(art)
24
    
25
# (optional) Log training metrics
26
wandb.log({"train_loss": 0.345, "val_loss": 0.456})
Copied!
The following code snippet leverages actively developed beta APIs and therefore is subject to change and not guaranteed to be backwards compatible. 
1
from wandb.beta.workflows import log_model
2
​
3
# (Optional) Declare an upstream dataset dependency
4
# see the `Declare Dataset Dependency` tab for
5
# alternative examples.
6
dataset = wandb.use_artifact("mnist:latest")
7
​
8
# This one mthod will serialize the model, start a run, create a version
9
# add the files to the version, and log the version. You can override
10
# the default name, project, aliases, metadata, and more!
11
log_model(model, "mnist-nn", aliases=["best"] if model_is_best else [])
12
​
13
# (optional) Log training metrics
14
wandb.log({"train_loss": 0.345, "val_loss": 0.456})
Copied!
Note: you may want to define custom serialization and deserialization strategies. You can do so by subclassing the _SavedModel class, similar to the _PytorchSavedModel class. All subclasses will automatically be loaded into the serialization registry. As this is a beta feature, please reach out to [email protected] with questions or comments.
If you would like to track your training data, you can declare a dependency by calling wandb.use_artifact on your dataset. Here are 3 examples of how you can declare a dataset dependency:
​
Dataset stored in W&B
1
dataset = wandb.use_artifact("[[entity/]project/]name:alias")
Copied!
​
Dataset stored on Local Filesystem
1
art = wandb.Artifact("dataset_name", "dataset")
2
art.add_dir("path/to/data") # or art.add_file("path/to/data.csv")
3
dataset = wandb.use_artifact(art)
Copied!
​
Dataset stored on Remote Bucket
1
art = wandb.Artifact("dataset_name", "dataset")
2
art.add_reference("s3://path/to/data")
3
dataset = wandb.use_artifact(art)
Copied!
After logging 1 or more Model Versions, you will notice that your will have a new Model Artifact in your Artifact Browser. Here, we can see the results of logging 5 versions to an artifact named mnist_nn-fz3l7wpe. 
If you are following along the example notebook, you should see a Run Workspace with charts similar to the image below
3. Link Model Versions to the Portfolio
Now, let's say that we are ready to link one of our Model Versions to the Model Portfolio. We can accomplish this manually as well as via an API.
Manual Linking
Programatic Linking
(Beta) Using `link_model`
The following video below demonstrates how to manually link a Model Version to your newly created Portfolio:
1.Navigate to the Model Version of interest
2.Click the link icon
3.Select the target Portfolio
4.(optional): Add additional aliases
While manual linking is useful for one-off Models, it is often useful to programmatically link Model Versions to a Portfolio - consider a nightly job or CI pipeline that wants to link the best Model Version from every training job. Depending on your context and use case, you may use one of 3 different linking APIs:
​
Fetch Model Artifact from Public API:
1
import wandb
2
​
3
# Fetch the Model Version via API
4
art = wandb.Api().artifact(...)
5
​
6
# Link the Model Version to the Model Portfolio
7
art.link("[[entity/]project/]portfolioName")
Copied!
​
Model Artifact is "used" by the current Run:
1
import wandb
2
​
3
# Initialize a W&B run to start tracking
4
wandb.init()
5
​
6
# Obtain a reference to a Model Version
7
art = wandb.use_artifact(...)
8
​
9
# Link the Model Version to the Model Portfolio
10
art.link("[[entity/]project/]portfolioName")
Copied!
​
Model Artifact is logged by the current Run:
1
import wandb
2
​
3
# Initialize a W&B run to start tracking
4
wandb.init()
5
​
6
# Create an Model Version
7
art = wandb.Artifact(...)
8
​
9
# Log the Model Version
10
wandb.log_artifact(art)
11
​
12
# Link the Model Version to the Portfolio
13
wandb.run.link_artifact(art, "[[entity/]project/]portfolioName")
Copied!
The following code snippet leverages actively developed beta APIs and therefore is subject to change and not guaranteed to be backwards compatible. 
In the case that you logged a model with the beta log_model discussed above, then you can use it's companion method: link_model
1
from wandb.beta.workflows import log_model, link_model
2
​
3
# Obtain a Model Version
4
model_version = wb.log_model(model, "mnist_nn")
5
​
6
# Link the Model Version
7
link_model(model_version, "[[entity/]project/]portfolioName")
Copied!
After you link the Model Version, you will see hyperlinks connecting the Version in the Portfolio to the source Artifact and visa versa.
4. Use a Model Version
Now we are ready to consume a Model - perhaps to evaluate its performance, make predictions against a dataset, or use in a live production context. Similar to logging a Model, you may choose to use the raw Artifact API or the more opinionated beta APIs.
Using Artifacts
(Beta) Using `use_model`
You can load in a Model Version using the use_artifact method.
1
import wandb
2
​
3
# Always initialize a W&B run to start tracking
4
wandb.init()
5
​
6
# Download your Model Version files
7
path = wandb.use_artifact("[[entity/]project/]portfolioName:latest").download()
8
​
9
# Deserialize your model (this will be your custom code to load in
10
# a model from disk - reversing the serialization process used in step 2)
11
model = make_model_from_data(path)
Copied!
The following code snippet leverages actively developed beta APIs and therefore is subject to change and not guaranteed to be backwards compatible. 
Directly manipulating model files and handling deserialization can be tricky - especially if you were not the one who serialized the model. As a companion to log_model, use_model automatically deserializes and reconstructs your model for use.
1
from wandb.beta.workflows import use_model
2
​
3
model = use_model("[[entity/]project/]portfolioName").model_obj()
Copied!
5. Evaluate Model Performance
After training many Models, you will likely want to evaluate the performance of those models. In most circumstances you will have some held-out data which serves as a test dataset, independent of the dataset your models have access to during training. To evaluate a Model Version, you will want to first complete step 4 above to load a model into memory. Then:
1.(Optional) Declare a data dependency to your evaluation data
2.Log metrics, media, tables, and anything else useful for evaluation
1
# ... continuation from 4
2
​
3
# (Optional) Declare an upstream evaluation dataset dependency
4
dataset = wandb.use_artifact("mnist-evaluation:latest")
5
​
6
# Evaluate your model in whatever way makes sense for your
7
loss, accuracy, predictions = evaluate_model(model, dataset)
8
​
9
# Log out metrics, images, tables, or any data useful for evaluation.
10
wandb.log({"loss": loss, "accuracy": accuracy, "predictions": predictions})
Copied!
If you are executing similar code, as demonstrated in the notebook, you should see a workspace similar to the image below - here we even show model predictions against the test data!
6. Promote a Version to Production
Next, you will likely want to denote which version in the Portfolio is intended to be used for Production. Here, we use the concept of aliases. Each Portfolio can have any aliases which make sense for your use case - however we often see production as the most common alias. Each alias can only be assigned to a single Version at a time.
via API
via UI Interface
Follow steps in Part 3. Link Model Versions to the Portfolio and add the aliases you want to the aliases parameter.
Coming soon!
The image below shows the new production alias added to v1 of the Portfolio!
7. Consume the Production Model
Finally, you will likely want to use your production Model for inference. To do so, simply follow the steps outlined in Part 4. Using a Model Version, with the production alias. For example:
1
wandb.use_artifact("[[entity/]project/]portfolioName:production")
Copied!
You can reference a Version within the Portfolio using different alias strategies:
latest - which will fetch the most recently linked Version
v# - using v0, v1, v2, ... you can fetch a specific version in the portfolio
production - you can use any custom alias that you and your team have assigned
8. Build a Reporting Dashboard
Use Weave to Build a Model Report