cancel
Showing results for 
Search instead for 
Did you mean: 
Machine Learning
Dive into the world of machine learning on the Databricks platform. Explore discussions on algorithms, model training, deployment, and more. Connect with ML enthusiasts and experts.
cancel
Showing results for 
Search instead for 
Did you mean: 

MLFlow: How to load results from model and continue training

Tilo
New Contributor

I'd like to continue / finetune training of an existing keras/tensorflow model.

We use MLFlow to store the model. How can I load the wieght from an existing model to the model and continue "fit" preferable with a different learning rate.

Just loading the model fails with some error, PyFuncModel does not provide the attributes and methods.

1 ACCEPTED SOLUTION

Accepted Solutions

Anonymous
Not applicable

@Tilo Wünsche​ 

To continue training an existing Keras/TensorFlow model that is stored in MLFlow, you need to follow the steps below:

  1. Load the model from MLFlow using mlflow.keras.load_model method.
import mlflow.keras
 
model = mlflow.keras.load_model("model_uri")
  1. Freeze the layers of the loaded model that you don't want to retrain.

for layer in model.layers[:-5]:
    layer.trainable = False

In this example, the last five layers will be trainable and the rest of the layers will be frozen.

  1. Compile the model with the desired learning rate and optimizer.

from tensorflow.keras.optimizers import Adam
 
optimizer = Adam(lr=0.0001)
model.compile(loss="categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])

In this example, the learning rate is set to 0.0001 and the Adam optimizer is used.

  1. Continue training the model with the new data. Use the fit method to continue training the model.
model.fit(new_X_train, new_y_train, epochs=10, batch_size=32, validation_data=(X_val, y_val))

In this example, the model is trained for 10 epochs with a batch size of 32.

With these steps, you should be able to load an existing Keras/TensorFlow model stored in MLFlow and continue training it with a different learning rate.

View solution in original post

3 REPLIES 3

Anonymous
Not applicable

@Tilo Wünsche​ 

To continue training an existing Keras/TensorFlow model that is stored in MLFlow, you need to follow the steps below:

  1. Load the model from MLFlow using mlflow.keras.load_model method.
import mlflow.keras
 
model = mlflow.keras.load_model("model_uri")
  1. Freeze the layers of the loaded model that you don't want to retrain.

for layer in model.layers[:-5]:
    layer.trainable = False

In this example, the last five layers will be trainable and the rest of the layers will be frozen.

  1. Compile the model with the desired learning rate and optimizer.

from tensorflow.keras.optimizers import Adam
 
optimizer = Adam(lr=0.0001)
model.compile(loss="categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])

In this example, the learning rate is set to 0.0001 and the Adam optimizer is used.

  1. Continue training the model with the new data. Use the fit method to continue training the model.
model.fit(new_X_train, new_y_train, epochs=10, batch_size=32, validation_data=(X_val, y_val))

In this example, the model is trained for 10 epochs with a batch size of 32.

With these steps, you should be able to load an existing Keras/TensorFlow model stored in MLFlow and continue training it with a different learning rate.

Orianh
Valued Contributor II

Basically you covered all what is needed for retrain model.

One thing that is good to mention is that if @Tilo Wünsche​  use DBR runtime 12.x .mlflow updated to version 2.x which throw an error if using mlflow.keras. ( No moudle names mlflow.keras)

from my point of view all remain the same, just use mlflow.tensorflow.load_model instead

Anonymous
Not applicable

Hi @Tilo Wünsche​ 

Hope all is well! Just wanted to check in if you were able to resolve your issue and would you be happy to share the solution or mark an answer as best? Else please let us know if you need more help. 

We'd love to hear from you.

Thanks!

Join 100K+ Data Experts: Register Now & Grow with Us!

Excited to expand your horizons with us? Click here to Register and begin your journey to success!

Already a member? Login and join your local regional user group! If there isn’t one near you, fill out this form and we’ll create one for you to join!