[Python] Is using mlflow.evaluate best practice for model assessment in Databricks?

I'm trying to assess the performance of a random forest regression model in Databricks. I'm familiar with Scikit-learn but new to Databricks and MLflow, so I'm not sure if I'm doing it in the best way.

This is the code I'm using to train, test, and log the model. I'll build on this in future by logging artifacts for things like feature importances.

with mlflow.start_run() as run:

#Train the model
model = RandomForestRegressor(random_state = 1)
model.fit(X = x_train, y = y_train)

#Use the model to make some test predictions
train_prediction = model.predict(X = x_train)
test_prediction = model.predict(X = x_test)

#Calculate metrics
test_mae = mean_absolute_error(y_true = y_test, y_pred = test_prediction)
test_mse = mean_squared_error(y_true = y_test, y_pred = test_prediction)
test_r2 = r2_score(y_true = y_test, y_pred = test_prediction)
test_rmse = np.sqrt(test_mse)

#Log metrics to the current run
mlflow.log_metrics({"test_mae": test_mae, "test_mse": test_mse, "test_r2": test_r2, "test_rmse": test_rmse})

#Log the model to the current run
signature = infer_signature(x_train, train_prediction)
mlflow.sklearn.log_model(sk_model = model, artifact_path = "my_model", signature = signature)


It works fine. However, all the MLflow docs I've seen do things differently, more like this:

#Train the model
model = RandomForestRegressor(random_state = 1)
model.fit(X = x_train, y = y_train)

#Infer the model's signature
train_prediction = model.predict(X = x_train)
signature = infer_signature(x_train, train_prediction)

#Construct evaluation dataset
evaluation_df = x_test
evaluation_df["target"] = y_test

#Log the model
with mlflow.start_run() as run:
model_info = mlflow.sklearn.log_model(sk_model = model, artifact_path = "my_model", signature = signature)
result = mlflow.evaluate(
model_info.model_uri,
evaluation_df,
targets = "target",
model_type = "regressor",
evaluators = ["default"],
)


I don't understand from reading the docs what advantages this gives. Am I missing out by not using mlflow.evaluate? Or are both methods equally valid ways to get the same result?

Any advice is much appreciated!

Continue reading...