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