Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Example of Saving an MLFlow
from sklearn.linear_model import ElasticNet
# these are internal wrapper/utility classes that we have developed to streamline the ML lifecycle process
from hs_mllib.model_lifecycle.packaging import MLModel, ScikitLearnModel
# this context MLFlow context manager allows experiment runs (parameters and metrics) to be tracked and easily queryable
with MLModel.mlflow.start_run() as run:
# data transformations and feature pre-processing code omitted (boiler-plate code)
...
# model construction
lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
# training
lr.fit(train_x, train_y)
# evaluate the model performance
predicted_qualities = lr.predict(test_x)
(rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)
# Wrap the model in our custom wrapper class
model = ScikitLearnModel(lr)
model.log_params(...)
model.log_metrics(...) # record the results of the run in ML Tracking Server
# optionally save model artifacts to object store and register model (give it a semantic version)
# so it can be built into a Sagemaker-servable Docker image
model.save(register=True)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment