Save MLFlow model in training job

lab08_mfl.py

@@ -46,11 +46,13 @@ def prepare_model(train_size_param, test_size_param, epochs, batch_size):
     )
     y_pred = model.predict(x_test.values)
 
+    input_example = X_test.values[10]
+
     rmse = mean_squared_error(y_test, y_pred)
 
     model.save("model_movies")
 
-    return model, rmse
+    return model, rmse, X_train, input_example
 
 
 train_size_param = float(sys.argv[1]) if len(sys.argv) > 1 else 0.8
@@ -66,7 +68,7 @@ with mlflow.start_run():
     mlflow.log_param("epochs", epochs)
     mlflow.log_param("batch size", batch_size)
 
-    model, rmse = prepare_model(
+    model, rmse, X_train, input_example = prepare_model(
         train_size_param=train_size_param,
         test_size_param=test_size_param,
         epochs=epochs,
@@ -75,4 +77,5 @@ with mlflow.start_run():
 
     mlflow.log_metric("RMSE", rmse)
 
-    mlflow.keras.log_model(model, "movies_imdb")
+    signature = mlflow.models.signature.infer_signature(X_train.values, model.predict(X_train.values))
+    mlflow.keras.save_model(model, "movies_imdb", input_example=input_example, signature=signature)

movies_imdb/MLmodel

@@ -0,0 +1,20 @@
+flavors:
+  keras:
+    data: data
+    keras_module: tensorflow.keras
+    keras_version: 2.4.0
+    save_format: tf
+  python_function:
+    data: data
+    env: conda.yaml
+    loader_module: mlflow.keras
+    python_version: 3.8.5
+saved_input_example_info:
+  artifact_path: input_example.json
+  format: tf-serving
+  type: ndarray
+signature:
+  inputs: '[{"type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1, 3]}}]'
+  outputs: '[{"type": "tensor", "tensor-spec": {"dtype": "float32", "shape": [-1,
+    1]}}]'
+utc_time_created: '2021-05-23 10:26:24.815501'

movies_imdb/conda.yaml

@@ -0,0 +1,10 @@
+channels:
+- defaults
+- conda-forge
+dependencies:
+- python=3.8.5
+- pip
+- pip:
+  - mlflow
+  - tensorflow==2.4.1
+name: mlflow-env

movies_imdb/data/keras_module.txt

@@ -0,0 +1 @@
+tensorflow.keras

movies_imdb/data/save_format.txt

@@ -0,0 +1 @@
+tf

movies_imdb/input_example.json

@@ -0,0 +1 @@
+{"inputs": [0.3076923076923066, 0.0018377625652137, 0.1073170731707316]}