Skip to main content

​
Log models

This guide shows you how to log model files to a W&B run as a model artifact, and how to retrieve and reuse those files in later runs. Logging models to a run lets you track which experiment produced a given model, capture its lineage, and make it available for downstream use. Use this page if you want lightweight model tracking as part of an experiment tracking workflow.
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 and prompts.
  • 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 about W&B Artifacts and advanced versioning use cases, see the Artifacts documentation.

​
Log a model to a run

This section explains how to use log_model to log a model artifact from a W&B run. Use the log_model method 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. For more information, see the Explore and traverse artifact graphs page within the Artifacts chapter. 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. Replace [YOUR-PROJECT] with your project name, [YOUR-ENTITY] with your entity name, [PATH-TO-MODEL] with the path to your model files, and [NAME] with a name for the model artifact. 
import wandb

# Initialize a W&B run
with wandb.init(project="[YOUR-PROJECT]", entity="[YOUR-ENTITY]") as run:

    # 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 isn’t specified, W&B uses 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 need the name of the model to retrieve the model path with the wandb.Run.use_model() method.
For parameters, see log_model in the API Reference. After the run completes, the model files are stored as a model artifact in W&B and are linked to the run as an output.

​
Download and use a logged model

After you log a model, you can retrieve its files in a later run to evaluate it, fine-tune it, or use it for inference. Use the use_model function to access and download model files previously logged to a W&B run. Provide the name of the model artifact where the model files you want to retrieve are stored. The name you provide must match the name of an existing logged model artifact. If you didn’t define name when you originally logged the files with log_model, the default name assigned is the basename of the input path, prepended with the run ID. Replace [YOUR-PROJECT] with your project name, [YOUR-ENTITY] with your entity name, and [YOUR-MODEL-NAME] with the name of the model artifact to retrieve. 
import wandb

# Initialize a run
with wandb.init(project="[YOUR-PROJECT]", entity="[YOUR-ENTITY]") as run:

    # 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 previous code snippet, the returned path is stored in a variable called downloaded_model_path. For parameters and return type, see use_model in the API Reference.