Skip to main content

Try in Colab

​
Log models

The following guide describes how to log models to a W&B run and interact with them.
The following APIs are useful for tracking models as a part of your experiment tracking workflow. Use the APIs listed on this page to log models to a run, and to access metrics, tables, media, and other objects.W&B suggests that you use W&B Artifacts if you want to:
  • Create and keep track of different versions of serialized data besides models, such as datasets, prompts, and more.
  • Explore lineage graphs of a model or any other objects tracked in W&B.
  • Interact with the model artifacts these methods created, such as updating properties (metadata, aliases, and descriptions)
For more information on W&B Artifacts and advanced versioning use cases, see the Artifacts documentation.

​
Log a model to a run

Use the log_model to log a model artifact that contains content within a directory you specify. The log_model method also marks the resulting model artifact as an output of the W&B run. You can track a model’s dependencies and the model’s associations if you mark the model as the input or output of a W&B run. View the lineage of the model within the W&B App UI. See the Explore and traverse artifact graphs page within the Artifacts chapter for more information. Provide the path where your model files are saved to the path parameter. The path can be a local file, directory, or reference URI to an external bucket such as s3://bucket/path. Ensure to replace values enclosed in <> with your own. 
import wandb

# Initialize a W&B run
run = wandb.init(project="<your-project>", entity="<your-entity>")

# Log the model
run.log_model(path="<path-to-model>", name="<name>")
Optionally provide a name for the model artifact for the name parameter. If name is not specified, W&B will use the basename of the input path prepended with the run ID as the name.
Keep track of the name that you, or W&B assigns, to the model. You will need the name of the model to retrieve the model path with the use_model method.
See log_model in the API Reference for parameters.

​
Download and use a logged model

Use the use_model function to access and download models files previously logged to a W&B run. Provide the name of the model artifact where the model files you are want to retrieve are stored. The name you provide must match the name of an existing logged model artifact. If you did not define name when originally logged the files with log_model, the default name assigned is the basename of the input path, prepended with the run ID. Ensure to replace other the values enclosed in <> with your own: 
import wandb

# Initialize a run
run = wandb.init(project="<your-project>", entity="<your-entity>")

# Access and download model. Returns path to downloaded artifact
downloaded_model_path = run.use_model(name="<your-model-name>")
The use_model function returns the path of downloaded model files. Keep track of this path if you want to link this model later. In the preceding code snippet, the returned path is stored in a variable called downloaded_model_path. See use_model in the API Reference for parameters and return type.
The link_model method is currently only compatible with the legacy W&B Model Registry, which will soon be deprecated. To learn how to link a model artifact to the new edition of model registry, visit the Registry linking guide.
Use the link_model method to log model files to a W&B Run and link it to the W&B Model Registry. If no registered model exists, W&B will create a new one for you with the name you provide for the registered_model_name parameter. Linking a model is analogous to ‘bookmarking’ or ‘publishing’ a model to a centralized team repository of models that others members of your team can view and consume. When you link a model, that model is not duplicated in the Registry or moved out of the project and into the registry. A linked model is a pointer to the original model in your project. Use the Registry to 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. A Registered Model is a collection or folder of linked model versions in the Model Registry. Registered models typically represent candidate models for a single modeling use case or task. The proceeding code snippet shows how to link a model with the link_model API. Ensure to replace other the values enclosed in <> with your own: 
import wandb

run = wandb.init(entity="<your-entity>", project="<your-project>")
run.link_model(path="<path-to-model>", registered_model_name="<registered-model-name>")
run.finish()
See link_model in the API Reference guide for optional parameters. If the registered-model-name matches the name of a registered model that already exists within the Model Registry, the model will be linked to that registered model. If no such registered model exists, a new one will be created and the model will be the first one linked. For example, suppose you have an existing registered model named “Fine-Tuned-Review-Autocompletion” in your Model Registry (see example here). And suppose that a few model versions are already linked to it: v0, v1, v2. If you call link_model with registered-model-name="Fine-Tuned-Review-Autocompletion", the new model will be linked to this existing registered model as v3. If no registered model with this name exists, a new one will be created and the new model will be linked as v0.
⌘I