MLFlow: How to load results from model and continue training

[Tilo]

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.

[Anonymous · Accepted solution]

@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]

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]

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!