search_mlflow_runs.py

import mlflow


EXPERIMENT_NAME = "fashionmnist-pytorch-demo"


def main():
    # Optional if using a local tracking server:
    mlflow.set_tracking_uri("http://localhost:5000")

    print("\nAll runs in experiment:")
    runs_df = mlflow.search_runs(
        experiment_names=[EXPERIMENT_NAME],
        order_by=["metrics.best_val_accuracy DESC"],
    )
    print(runs_df[[
        "run_id",
        "tags.mlflow.runName",
        "tags.group_name",
        "tags.model_name",
        "params.learning_rate",
        "params.batch_size",
        "metrics.best_val_accuracy",
        "metrics.val_accuracy",
    ]])

    print("\nOnly model_comparison runs:")
    model_cmp_df = mlflow.search_runs(
        experiment_names=[EXPERIMENT_NAME],
        filter_string='tags.group_name = "model_comparison"',
        order_by=["metrics.best_val_accuracy DESC"],
    )
    print(model_cmp_df[[
        "run_id",
        "tags.mlflow.runName",
        "tags.model_name",
        "metrics.best_val_accuracy",
    ]])

    print("\nOnly CNN runs with strong accuracy:")
    cnn_df = mlflow.search_runs(
        experiment_names=[EXPERIMENT_NAME],
        filter_string='tags.model_name = "cnn" AND metrics.best_val_accuracy > 0.85',
        order_by=["metrics.best_val_accuracy DESC"],
    )
    print(cnn_df[[
        "run_id",
        "tags.mlflow.runName",
        "metrics.best_val_accuracy",
    ]])


if __name__ == "__main__":
    main()