I am logging the model using
mlflow.sklearn.log_model(model, "my-model")
and I want to set tags to the model during logging, I checked that this method does not allow to set tags, there is a mlflow.set_tags()
method but it is tagging the run not the model.
Does anyone know how to tag the model during logging?
Thank you!
When using mlflow.sklearn.log_model
you work with the experiment registry which is run-focused so only experiments and runs can be described and tagged.
If you want to set tags on models, you need to work with the model registry.
The solution I would recommend is to register the model when logging using registered_model_name
(there are more fine-grained ways, too) and use MLFlowClient
API to set custom properties (like tags) of the already registered model.
Here is a working example:
import mlflow
from mlflow.client import MlflowClient
mlflow.set_tracking_uri('http://0.0.0.0:5000')
experiment_name = 'test_mlflow'
try:
experiment_id = mlflow.create_experiment(experiment_name)
except:
experiment_id = mlflow.get_experiment_by_name(experiment_name).experiment_id
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score
with mlflow.start_run(experiment_id = experiment_id):
# log performance and register the model
X, y = load_iris(return_X_y=True)
params = {"C": 0.1, "random_state": 42}
mlflow.log_params(params)
lr = LogisticRegression(**params).fit(X, y)
y_pred = lr.predict(X)
mlflow.log_metric("accuracy", accuracy_score(y, y_pred))
mlflow.sklearn.log_model(lr,
artifact_path="models",
registered_model_name='test-model'
)
# set extra tags on the model
client = MlflowClient(mlflow.get_tracking_uri())
model_info = client.get_latest_versions('test-model')[0]
client.set_model_version_tag(
name='test-model',
version=model_info.version,
key='task',
value='regression'
)
Here is the illustration
See also this excellent documentation of MLFlow Client.