Logging and Registering Models in W&B

It's never been easier to log your model checkpoints, keep track of the best ones, and maintain lineage of runs and results!

W&B is introducing a few convenience methods to make logging models and linking them to the registry simple:

• log_model
• use_model
• link_model

Imports

Log in to W&B

• You can explicitly login using wandb login or wandb.login() (See below)
• Alternatively you can set environment variables. There are several env variables which you can set to change the behavior of W&B logging. The most important are:
  • WANDB_API_KEY - create a new API key in your "Settings" section under your profile at wandb.ai/settings
  • WANDB_BASE_URL - this is the url of the W&B server
• Create a new API key in "Profile" -> "Settings" in the W&B App. Store your API key securely. It can only be viewed once when created.

Define the Model and Dataset

This is a simple implementation of a Vision Transformer (ViT) and utilizes a random dataset for training.

• Credit to https://github.com/lucidrains/vit-pytorch

Define some config for the model and dataset

Log Model Checkpoints to W&B with 1 Line!

Use the log_model method to log a model artifact containing the contents inside the ‘path’ to an a run. It also marks it as an output to the run. You can see the full lineage graph of the model artifact by accessing the lineage tab inside the Artifacts view.

log_model() takes as input:

• path: A path to the model file(s), which can be a local file (of the form /local/directory/file.txt), directory (of the form /local/directory), or reference path to S3 (s3://bucket/path).
• name: An optional name for the model artifact the files will be logged to. Note that if no name is specified, This will default to the basename of the input path prepended with the run ID.
• aliases: An optional list of aliases, which can be thought of as semantic ‘nicknames’ or identifiers for a model version. For example, if this model yielded the best accuracy, you might add the alias ‘highest-accuracy’ or ‘best’.

Link Your Best Models to the Model Registry

You can bookmark your best model checkpoints and centralize them across your team. The Model Registry allows you can organize your best models by task, manage model lifecycle, facilitate easy tracking and auditing throughout the ML lifecyle, and automate downstream actions with webhooks or jobs. You can this via api through link_model(), which takes as input:

• path: A path to the model file(s), which can be a local file (of the form /local/directory/file.txt), directory (of the form /local/directory), or reference path to S3 (s3://bucket/path).
• registered_model_name: the name of the Registered Model - a collection of linked model versions in the Model Registry, typically representing a team’s ML task - that the model should be linked to. If no Registered Model with the given name exists, a new one will be created with this name.
• name: An optional name for the model artifact the files will be logged to. Note that if no name is specified, This will default to the basename of the input path prepended with the run ID.
• aliases: An optional list of aliases, which can be thought of as semantic ‘nicknames’ or identifiers for a linked model version. For example, since this model is being linked, or published, to the Model Registry, you might add an alias “staging” or “QA”.