Skip to content

Log models

Comet's models allow you to log trained models for all the most popular machine learning frameworks.

If your framework has an existing integration with Comet, please double-check if Comet provides a specialized model logging method (e.g., with comet_ml.integration.pytorch.log_model()) which allows you to skip saving the model to disk as an intermediate step.

Example of models in the Assets & Artifacts tab
View model in the Comet UI for Comet Tutorial 2

Comet allows you to access, download, and register (for production) a logged model at any time from the Single Experiment page, and recommends you to log both model checkpoints and final model for maximum reproducibility.

The following method can be used to log model:

The models logged to an Experiment can be viewed in the following Single Experiment Tab:

For example, you could...

Resume training from any model checkpoint allowing for further fine-tuning or adaptation to new data without starting from scratch.

Warning

While it is possible to log models using Comet Artifacts, this is not the recommended approach. We recommend logging models to Experiments so that they can be promoted to the Model Registry if needed.

Log a model file

The example below showcases how to log a mock .txt model to Comet.

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

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

# Create an example model file
with open("/tmp/model.txt", "w+") as f:
    f.write("model")

# Log the model to Comet
exp.log_model(
    name="example-model-from-file",
    file_or_folder="/tmp/model.txt"
)

The file_or_folder argument support a single file path, folder paths, and file-like objects too. For single file paths and file-like objects, you can also specify a custom file name with the file_name argument; for a folder, you can specify the prepend_folder_name boolean argument to define whether to prepend the file path by the folder name.

Additionally, you could use the overwrite argument to overwrite any existing model with the same name or add a metadata argument to log extra information (such as tags) in a JSON format.

Log a model object

The example below showcases how to log a model from a file-like object.

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

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

# Create an example model object
model_object = io.BytesIO()
model_object.write("mode".encode())
model_object.seek(0)

# Log the model to Comet
exp.log_model(
    name="example-model-from-object",
    file_or_folder=model_object,
    file_name="model.txt",
)

Using a file-like object allows you to avoid saving the model to disk.

Please refer to Log a model file above for information on additional optional arguments.

Log a Pytorch model

The example below showcases how to log a Pytorch model with the comet_ml.integration.pytorch.log_model() method.

 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
import comet_ml
from comet_ml.integration.pytorch import log_model
import torch
import torch.nn as nn

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

# Create an example pytorch model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()

    def forward(self, x):
        return 1

pytorch_model = SimpleModel()

# Log the pytorch model to Comet
log_model(
    experiment=exp,
    model=pytorch_model,
    model_name="example-model-pytorch",
)

Additionally, you can specify metadata, pickle_module, and any kwargs supported by torch.save().

Log a Scikit-Learn model

The example below showcases how to log an sklearn model with the comet_ml.integration.sklearn.log_model() method.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
import comet_ml
from comet_ml.integration.sklearn import log_model
from sklearn.svm import SVC

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

# Create an example sklearn model
sklearn_model = SVC()

# Log the sklearn model to Comet
log_model(
    experiment=exp,
    model_name="example-model-sklearn",
    model=sklearn_model,
)

Additionally, you can specify metadata, persistence_module, and any dump kwargs supported by sklearn.

Apr. 29, 2024