mlflow_sklearn_autolog_random_forest_diabetes_quickstart.py

1import mlflow
2from sklearn.model_selection import train_test_split
3from sklearn.datasets import load_diabetes
4from sklearn.ensemble import RandomForestRegressor
5
6# Load the diabetes dataset
7db = load_diabetes()
8X_train, X_test, y_train, y_test = train_test_split(db.data, db.target)
9
10# Connect to the MLflow tracking server (local default)
11mlflow.set_tracking_uri("http://127.0.0.1:5000")
12
13# Sets the current active experiment to the "Quickstart" experiment and
14# returns the Experiment metadata
15mlflow.set_experiment("MLflow Quickstart")
16
17# Enable autologging
18mlflow.sklearn.autolog()
19
20# Create and train a model
21rf = RandomForestRegressor(n_estimators=100, max_depth=6, max_features=3)
22rf.fit(X_train, y_train)
23
24# Use the model to make predictions on the test dataset
25predictions = rf.predict(X_test)