Weights & Biases Registry Tutorial

The goal of this notebook is to demonstrate how you and other members of your organization can use W&B Registry to track, share, and use dataset and model artifacts in your machine learning workflows. By the end of this notebook, you will know how to use W&B to:

1. Create collections within W&B Registry
2. Make dataset and model artifacts available to other members of your organization, and
3. Download your trained model and dataset artifacts from the registry for inference

To achieve this, we will train a neural network to identify animal classes (mammal, amphibian, reptile, and so forth) based on features such as weather or not they ahve feathers, fins, and so on.

Install and import packages

Retrieve and process dataset

We will use the open source Zoo dataset from the UCI Machine Learning Repository.

Retrieve dataset

We can either manually download the dataset or use the ucimlrepo package to import the dataset directly into our notebook. For this example, we will import the dataset directly into this notebook:

Explore the data

Let's take a quick look at the shape and data type of the dataset:

Process data

For training let's convert our dataset from a pandas DataFrame to a tensor with PyTorch, convert the data type of our input tensor(float64 to float32) to match the data type of the nn.Linear module, and convert our label tensor to index from 0-6:

Save processed dataset locally using torch.save

Track and publish dataset

Within the Dataset registry we will create a collection called "zoo-dataset-tensors". A collection is a set of linked artifact versions in a registry.

To create a collection we need to do two things:

1. Specify the collection and registry we want to link our artifact version to. To do this, we specify a "target path" for our artifact version.
2. Use the wandb.run.link_artifact method and pass our artifact object and the target path.

Define target path of the collection

The target path of a collection consists of two parts:

• The name of the registry
• The name of the collection within the registry

If you know these two fields, you can create the full name yourself with string concatenation, f-strings, and so forth:

target_path = f"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}"

Publish dataset to registry

Let's publish our dataset to the Dataset registry in a collection called "zoo-dataset-tensors". To do this, we will

1. Get or create the target path. (For this notebook we will programmatically create the target path).
2. Initialize a run
3. Create an artifact object
4. Add each split dataset as individual files to the artifact object
5. Link the artifact object to the collection with run.link_artifact(). Here we specify the target path and the artifact we want to link.

First, let's create the target path:

Now that we have the target path, let's publish the dataset to the "Dataset" registry. In the following code cell, ensure to replace the values enclosed in <> with your team's entity:

Split data and publish split dataset to registry

Split the data into a training and test set. Splitting the dataset and tracking them as separate files will make it easier for a different user to use the same datasets for future reproducibility, testing, and analysis.

Next, let's publish this dataset into a different collection within the Dataset registry called "zoo-dataset-tensors-split":

We can verify we correctly linked our artifact to our desired collection and registry with W&B App UI:

1. Navigate to the Registry App
2. Select on the Dataset registry
3. Click View details "zoo-dataset-tensors-split" collection
4. Click the View button next to the artifact version
5. Select the Files tab

You should see four files: "zoo_dataset_X_test", "zoo_dataset_X_train", "zoo_labels_y_test", and "zoo_labels_y_train".

Define a model

The following cells show how to create a simple neural network classifier with PyTorch. There is nothing unique about this model, so we'll will not go into detail of this code block.

Define hyperparameters, loss function, and optimizer

Train model

Next, let's train a model using the training data we published to the registry earlier in this notebook. After we train the model, we will publish that model to W&B.

To do this, let's first get the artifact we published to the "Dataset" registry. To retrieve an artifact from a registry, we need to know the name of that artifact. The name of an artifact in a registry consists of the prefix wandb-registry-, the name of the registry, the name of the collection, and the artifact version:

# Artifact name/filepath for downloading and using artifacts published to a registry
f"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}:v{VERSION}"

Since we only linked one artifact version, the version we'll use is v0. (W&B uses 0 indexing).

Note that the name of an artifact is nearly identical to the target path we specified in a previous step when we publish our artifact to the registry except for the version number:

# Target path for publising an artifact version to a registry
f"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}"

The preceding cell might look intimidating. Let's break it down:

• First, we download the dataset from the Dataset registry and load it as a tensor
• Next, we create a simple training loop
  • Within the training loop we log the loss for each step
  • We checkpoint(save) the model every time the remainder of the epoch divided by 100 is 0 and the loss is lower than the previously recorded loss.
  • We then add the saved PyTorch model to our artifact object.

A couple of things to note:

1. The preceding code cell adds a single artifact version to W&B. You can confirm this by navigating to your project workspace, select Artifacts in the left navigation, and under models click the name of the artifact (starts with zoo-{run.id}). You will see a single model with version v0.
2. At this point, we have only tracked the model artifact within our team's project. Anyone outside of our team does not have access to the model we created. To make this model accessible to members outside of our team, we will need to publish our model to the registry.

Publish model to registry

Let's make this model artifact available to other users in our organization. To do this, we will create a collection within the Model registry.

To create a collection within a registry, we need to know the full name of the artifact. The full name of the artifact consists of the name we provided to it when we created the Artifact object and its location within our team's project.

There are two ways to get the full name of an artifact, interatively with the W&B App UI or programmatically with the W&B Python SDK. In this example, we'll programmatically create the name of the artifact since we have these values loaded into memory.

Programmatically create name of artifact

The full name of an artifact consists of four components:

• Team entity
• Project name
• The name of the artifact (the string you passed when you created the artifact object with wandb.Artifact())
• The artifact version

Putting this together, the full name of an artifact is:

# Full name of an artifact in a team project
artifact_name = f'{TEAM_ENTITY}/{PROJECT}/{ARTIFACT_NAME}:v{VERSION}'

Now that we have the full name of our model artifact. Let's publish it to the model registry.

Similar to how we created a target path when we published our dataset artifact to the Dataset registry, let's create the target path for our model artifact. The target path tells W&B the collection and registry (Model registry) to link our artifact version to.

As a reminder, the target path to link an artifact to a registry consists of:

# Target path used to link artifact to registry
target_path = f"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}"

Putting this all together, we specify our artifact name in run.use_artifact() and the target path for run.link_artifact():

The preceding code block links our model artifact version to a collection called "Zoo_Classifier_Models" within the model registry.

View lineage map of registered model

Let's say that you did not know exactly which model version to use. You can check the lineage of all artifact versions on the W&B App UI. The lineage shows which artifacts were used as input to a run and which artifacts were the output of a given run.

For example, the image below shows the "Zoo_Classifier_Models" collection within the model registry. Highlighted in yellow is the current model artifact version that is linked to the registry.

From left to right we see that the run "trim-rain-2" was responible for creating the "split_zoo_dataset" artifact. (Recall that this is the dataset artifact that contains the test and training data).

We then see that the "golden-sunset-3" run used the "split_zoo_dataset" artifact for training. Within this run, we created a model artifact. The specific artifact version we linked to "Zoo_Classifier_Models" is called zoo-wyhak4o0:v10.

To view the lineage map of an artifact in a registry:

1. Navigate to the Registry app at https://wandb.ai/registry
2. Click on a registry
3. Select an artifact version
4. Select the Lineage tab

Download artifacts from registry for inference

For this last section, suppose you are a different user in a different team within the same organization. You and your team want to download the model and test dataset that was published to your organization's registry by a different team. You and your team will use the model and test dataset for inference and store those findings in a project called "Check_Zoo_Model".

Note: The team member that wants do use and download published artifacts has member role permissions. This means they can view and download artifacts from the registry.

How can you retrieve the artifacts version that were published by another team? Simple:

1. Get the full name of the artifact version programmatically or interactively with the W&B App UI
2. Use the W&B Python SDK to download the artifacts

Interactively get full name of model and dataset artifacts from registry

1. Go to the W&B Registry app at https://wandb.ai/registry/.
2. Select the registry that your artifact is linked to.
3. Click the View details button next to the name of the collection with your linked artifact.
4. Click on the View button next to the artifact version.
5. Within the Version tab, copy path listed next to Full Name.
6. Paste the full name of the registry for the artifact_or_name field in run.use_artifact().

Note: In this example, we happen to know these values, so we'll programmatically create the full name of our model and dataset artifacts published in the registry.

Get trained model from model registry

In the following code cell, ensure to replace the values enclosed in <> with the entity of a different team in your organization than the one you specified earlier in this notebook.

Note: If you do not have another team entity, you can re-use the entity you specified earlier.

For PyTorch models, we need to redefine our model architecture:

Get test dataset from Dataset registry

Let's get the test dataset from our registry. Similar to the above code block, we will specify the full name of the artifact version we want from our Dataset registry.

Make predictions with loaded model

How does our model perform? Recall that the goal of the neural network is to predict the animall class based on features of that animal.

For each prediction, our model returns an integer that refers the class. Let's create a dictionary so we can map the integer to the class name:

Let's feed our model some data to make predictions:

These integers don't mean much, let's convert them to return the animal class and store this into a pandas DataFrame for us to compare the predicted vs the true values:

Let's see how many it predicted corrected

Let's view these as percentages:

The percentage the model predicted correct might vary. As of writing this notebook, our the model correctly predicted ~88% of the time:

As a next step, you could dig into the examples that were incorrectly predicted and try to figure out why it predicted incorrectly. You could also try feature engineering to extract more features to train with.

Summary

In this notebook completed each major step in a typical machine learning workflow, from downloading a dataset, processing the dataset, defining a model, training that model on processed data, checking/saving the best model, and checking how the model performed by making predictions with that model on a dataset it had not seen before.

Throughout the process you learned how to use W&B Registry to:

• Track and publish multiple datasets
• Track and publish a model
• Mimic how someone else in your organization (with correct permission) can download and use model and datasets published in the W&B Registry for further analysis.

Next steps:

As the number of machine learning experiments increases, so does the complexity of keeping track of saved models and datasets. For each model version, we recommend that you document key aspects of your model such as a brief summary of the model, information about the architecture of the model, how someone can deserialize a saved model, and so forth. You can provide all of this information, and more, within the Description field of the model version.