Skip to content

Track your trained models

Comet Model Registry allows you to track and manage all your models in a single place.

Comet Model Registry - Table
The Model Registry page for the examples workspace

For each registered model, you can review relevant information such as model card, lineage, usage instructions, set tags, and model status. And, you can connect all your model activity via web-hooks to your CI/CD pipeline.

Get started with Model Registry

The Comet Model Registry is defined at the workspace level, but supports tracking trained models at the experiment level.

Given a Comet Experiment, the general steps to save a model to the Model Registry are:

  1. Log the trained model using the log_model() method offered by the Comet SDK.
  2. Register the model, either from the Comet UI or Comet SDK.
  3. Manage versions, history, and webhooks for the registered model.
  4. Deploy the registered model.

Below you can find a step-by-step walkthrough of an end-to-end example.

Prerequisites

Make sure that you have correctly configured the Comet SDK and that you have the code set ready to create and run the experiment for model training / tuning.

1. Log the trained model

You can log any trained or fine-tuned model to Comet, whether it is saved to local disk / cloud or hosted in memory as a file-like object.

The example below trains a dummy sklearn classifier and logs it from filepath.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import comet_ml
import joblib
from sklearn.datasets import make_classification
from sklearn.dummy import DummyClassifier
from sklearn.model_selection import train_test_split

# Initialize the Comet Experiment
comet_ml.init()
exp = comet_ml.Experiment()

# Define the dataset
X, y = make_classification(n_samples=100, n_features=4, n_classes=2, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Train the model
dummy_clf = DummyClassifier(strategy="most_frequent")
dummy_clf.fit(X, y)

# Save the model to local filepath
model_filepath = "dummy_classifier_model.joblib"
joblib.dump(dummy_clf, model_filepath)

# Log the model to Comet
exp.log_model(
    name="dummy classifier",
    file_or_folder=model_filepath,
    metadata={"framework": "sklearn"},
)

The logged model is now available from the Assets & Artifacts tab of the Single Experiment page of the Comet UI. The url for the page is printed in the terminal after experiment creation.

Comet Single Experiment page - Logged model
The example logged model accessed from the Single Experiment page

2. Register the model

When you click on the + Register button by the model's folder name in the Assets & Artifacts tab of the Single Experiment page, the Register Model windows pops-up.

Comet Single Experiment page - Register Model
The Register Model pop-up window

Note

If you have already registered models, the Register Model window will appear different to allow you the ability to either save to an existing registered model as a new version or to register a new model entirely.

When you register the model, you can specify:

  • Name: required; auto-populated by Comet with the logged model name.
  • Version: required; auto-populated by Comet with "1.0.0".
  • Description: optional.
  • Visibility: required; either public or private.

The Comet UI highlights registered models in the Assets & Artifacts tab by adding a View Existing button by the right side of any registered model asset as showcased in the screenshot below.

Comet Single Experiment page - View registered version for logged model
The registered model version viewed from the Single Experiment page

Comet supports other flows for registering the model, both from Comet UI and from Comet SDK. For more instructions, please refer to the Register a model docs page.

3. Manage the registered model

When you click on the View Existing button by a registered model's folder name in the Assets & Artifacts tab of the Single Experiment page, the Comet UI displays information about the registered model version including a link to the corresponding Model Registry page.

Alternatively, you can access the Model Registry page from your workspace home page by clicking on the Model Registry tabbed page.

Comet Model Registry - Model page
The model page for a registered model in Comet Model Registry

Model Registry supports a variety of management actions:

  • From the Model Versions page, you can:

    • Change the model status.
    • Manage tags.
    • Customize the table view by adding and removing columns.
    • ...

    Find out more from the Manage a registered model docs page.

  • From the Notes page, you can add free-form text information about the model useful for team collaboration.

  • From the History page, you can review all actions performed on the model across the days.
  • From the Webhooks page, you can create webhooks based on changes to the model version, tags, and/or status. Find out more from the Configure webhooks docs page.
  • From the Pending requests page, workspace owners can manage user requests to change model status.

4. Deploy the model

You can download and deploy your model programmatically via Comet UI or Comet SDK.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
import comet_ml

# Initialize the Comet SDK
comet_ml.init()

# Define a Comet API object
api = comet_ml.API()

# Get the Model object
model = api.get_model(
    workspace=api.get_default_workspace(), model_name="dummy classifier",
)

# Download the Model from Comet Model Registry
model.download("1.0.0", output_folder="./", expand=True)

Find more details in the Python SDK Model references.

To download a registered model as a .zip file, you can either:

• Click the download icon next to the model name in the Assets & Artifacts tab of the Single Experiment page.

Single Experiment page - Download model

• Download a specific version of the model through the table side menu of the Model Registry page.

Model Registry page - Download version

The downloaded .zip file contains all of the files that were logged with the Experiment model as described in step 1. Log the trained model.

Note

If you logged remote assets to your models that are stored in AWS S3 or GCP GCS, assuming you have the right credentials and permissions configured, Model.download() will download them as well.

Learn more

Jun. 14, 2024