mlflow_sklearn_logistic_regression_experiment_tracking_quickstart.py

1import mlflow
2from mlflow.models import infer_signature
3
4from sklearn(import datasets
5from sklearn.model_selection import train_test_split
6from sklearn.linear_model import LogisticRegression
7from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
8
9# Load the Iris dataset
10X, y = datasets.load_iris(return_X_y=True)
11
12# Split the data into training and test sets
13X_train, X_test, y_train, y_test = train_test_split(
14    X, y, test_size=0.2, random_state=42
15)
16
17# Define the model hyperparameters
18params = {
19    "solver": "lbfgs",
20    "max_iter": 1000,
21    "multi_class": "auto",
22    "random_state": 8888,
23}
24
25# Train the model
26lr = LogisticRegression(**params)
27lr.fit(X_train, y_train)
28
29# Predict on the test set
30y_pred = lr.predict(X_test)
31
32# Calculate metrics
33accuracy = accuracy_score(y_test, y_pred)
34
35# Set our tracking server uri for logging
36mlflow.set_tracking_uri(uri="http://127.0.0.1:8080")
37
38# Create a new MLflow Experiment
39mlflow.set_experiment("MLflow Quickstart")
40
41# Start an MLflow run
42with mlflow.start_run():
43    # Log the hyperparameters
44    mlflow.log_params(params)
45
46    # Log the loss metric
47    mlflow.log_metric("accuracy", accuracy)
48
49    # Set a tag that we can use to remind ourselves what this run was for
50    mlflow.set_tag("Training Info", "Basic LR model for iris data")
51
52    # Infer the model signature
53    signature = infer_signature(X_train, lr.predict(X_train))
54
55    # Log the model
56    model_info = mlflow.sklearn.log_model(
57        sk_model=lr,
58        artifact_path="iris_model",
59        signature=signature,
60        input_example=X_train,
61        registered_model_name="tracking-quickstart",
62    )