This is a common situation for data scientists, especially those with cross-collaborative teams, stakeholders, and changing priorities. Model development can quickly become a wild goose chase of figuring out requirements, data sources, performance metrics, etc. Through all of this chaos, it is critical that teams are able to not only keep track of all phases of model development, but are also able to reproduce results quickly. You need to model governance.
Model Governance
Model governance refers to the framework by which a team maintains controls around its use of models, including experiment tracking, version control, reproducibility, etc. It is a critical capability that many teams do not consider until their models are already in production. It is important to have a process to effectively track development efforts, models in production, and data leveraged so that organizations can ensure quality, compliance, and monitor machine learning integration.
To ensure that the right models are in production, we need a way to manage various models and versions, track performance metrics, and reproduce results. This is where ML Flow comes in. ML Flow is an open-source ML Ops platform that allows flexible model development, deployment, and management by streamlining logging and tracking of models, metrics, and more.
This article will give an introduction into how to install ML Flow, create experiments, log models and more.
ML Flow Tutorial
Model Development
For our model, we will construct a basic linear regression model. This is a model to predict the annual salary of data professionals using inputs that include job title, company size, etc.
See data here: https://www.kaggle.com/datasets/ruchi798/data-science-job-salaries (CC0: Public Domain). I slightly modified the data to reduce the number of options for certain features. Please note that this is a very poor model. We are using it only to demonstrate ML Flow.
#define independent and dependent features
X = salary_data.drop(columns = 'salary_in_usd')
y = salary_data['salary_in_usd']
#split between training and testing sets
X_train, X_test, y_train, y_test = train_test_split(
X, y, random_state = 104, test_size = 0.2, shuffle = True)
#define parameters separately to log
params = {"fit_intercept": True,
"positive": False}
#fit linear regression model
regr = linear_model.LinearRegression(**params)
regr.fit(X_train, y_train)
#make predictions
y_pred = regr.predict(X_test)Note that we define the model parameters separately. This will be important when we incorporate ML Flow.
Logging Model with ML Flow
To start using ML Flow, you can navigate to the terminal and pip install it as a module. Next, run the server command to host a local server on your machine and leverage the ML Flow user interface.
pip install mlflowmlflow server --host 127.0.0.1 --port 8080Now that we have ML Flow installed and our local server running, we can log our model. Before logging a model, we need to start a new ‘experiment’. Think of this as a project, a bin where we will log and store all related models, metrics, data. etc. In the code below, we start a new experiment titled “Logging Example”.
An experiment consists of runs, which contain the model used, its parameters, metrics, artifacts, and anything else that you choose to log. As you log more runs, we can track our progress and compare the performance of the various models developed.
Using the code below, we:
- Start the experiment
- Start a run titled “salary_baseline_regression”
- Log the adjusted R2 as a metric
- Log the model’s parameters
- Set a tag for the model
- Log the actual model
import mlflow
from mlflow.models import infer_signature
#set tracking uri
mlflow.set_tracking_uri(uri = "http://127.0.0.1:8080")
#create a new MLflow Experiment
mlflow.set_experiment("Logging Example")
#start an MLflow run
with mlflow.start_run(run_name = "salary_baseline_regression") as run:
# Log the loss metric
mlflow.log_metric("Adjusted R2", r2_score(y_test, y_pred))
#log the hyperparameters
mlflow.log_params(params)
#set a tag that we can use to remind ourselves what this run was for
mlflow.set_tag("Model Type", "Baseline")
#log the model
model_info = mlflow.sklearn.log_model(
sk_model = regr,
artifact_path = 'model',
signature = infer_signature(X_train, regr.predict(X_train)),
input_example = X_train,
registered_model_name = "salary_baseline_regression")After running, you can navigate to the ML Flow UI using the following link: http://localhost:8080/
Navigating ML Flow UI
The ML Flow UI will open on the ‘Experiments’ tab. Here, we can see the experiment we created, the run for our baseline model, the model used, as well as metrics and parameters we logged.
Note that the metrics and parameters may not be displayed at first. You can add them to the table using the ‘Columns’ drop down.
Next, click on the run name to navigate to the run view. This brings us to the overview tab where we can clearly see everything we have logged (model, R2 metric, model parameters).
Next, click on the model name or navigate to the model view. Here we can view our model, add descriptions and/or tags, and keep track of it’s versions. The version number will increase every time we create a run using this model.
The tags and descriptions are helpful for overall documentation and to tag models in development vs. under review vs. in production.
So far, we’ve seen how to create experiments and runs, as well as log models, metrics, parameters, etc. These are the very basics of ML Flow. Now let’s explore more options to automate some of the process and make best use out of ML Flow.
Advanced Logging Options
There are likely many metrics we would like to track, which will vary depending on the model and approach. Adding and removing metrics every time we change our approach can be laborious. However, we can feed ML Flow our test data and results to automate the logging of metrics and streamline our process.
First, let’s save a copy of our test data and the results of our model.
#merge salaries and predictions for mlfow
eval_data = X_test.copy()
eval_data["salary"] = y_testNext, we give ML Flow our data and tell it what our target variable is (salary). Now we can add a few things to our ML Flow process:
- Log our data to keep track of model inputs
- Use mlflow.evaluate() to automate logging the model metrics
We can also log ‘artifacts’, which can be many things. During development and evaluation, it may be helpful to develop visualizations that reflect model performance. We can log those as well using mlflow.log_artifact().
#set tracking uri
mlflow.set_tracking_uri(uri = "http://127.0.0.1:8080")
#create an instance of a PandasDataset
dataset = mlflow.data.from_pandas(
eval_data, source = 'kaggle', name = "ds_salaries", targets = "salary")
#create a new MLflow Experiment
mlflow.set_experiment("Salary Prediction")
#start an MLflow run
with mlflow.start_run(run_name = "salary_baseline_regression") as run:
#log the dataset
mlflow.log_input(dataset, context = "training")
#log the hyperparameters
mlflow.log_params(params)
#log pngs
mlflow.log_artifact('Result.png')
#set a tag that we can use to remind ourselves what this run was for
mlflow.set_tag("Model Type", "Baseline")
#log the model
model_info = mlflow.sklearn.log_model(
sk_model = regr,
artifact_path = 'model',
signature = infer_signature(X_train, regr.predict(X_train)),
input_example = X_train,
registered_model_name = "salary_baseline_regression")
model_uri = mlflow.get_artifact_uri("model")
#create the model uri and evalute the model
result = mlflow.evaluate(model = model_uri, data = dataset, model_type = "regressor", evaluators=["default"])Now, if we navigate back to the ML Flow UI and click on the run, we can see that there are several metrics we are tracking.
And if we want to evaluate multiple models in a run, we can compare their performance side-by-side using the automated visualizations below.
To view the visualizations you logged, navigate to the artifacts view.
Lastly, we can now view our our data inputs and their types. It is best to log datasets as we drop or add new features, sources, etc.
Recalling ML Flow Models
To reproduce a model that we have logged, we can easily recall models from ML Flow using the model name and version. This can also be used to retrain models that are currently in production.
model_name = "salary_baseline_regression"
model_version = "latest"
# Load the model from the Model Registry
model_uri = f"models:/{model_name}/{model_version}"
model = mlflow.sklearn.load_model(model_uri)Databricks Integration
If you are (like me) a fan of Databricks, ML Flow is fully integrated within the platform. No need to run ML Flow locally, as Databricks integrates the ML Flow UI within their ‘Experiments’ view. It also includes additional features like tracking for generative AI applications, agents, and integrates AutoML for model optimization. See full documentation here: https://docs.databricks.com/aws/en/mlflow/
Conclusion
In conclusion, model governance is a must-have capability for any organization to ensure quality and compliance. Tracking and monitoring tools like ML Flow offer solutions to streamline model governance processes and fully manage model development and deployment.
I hope you have enjoyed my article! Please feel free to comment, ask questions, or request other topics.
Connect with me on LinkedIn: https://www.linkedin.com/in/alexdavis2020/
