From 8c94666b251ceaf18b90e278b3ac48b7cb836178 Mon Sep 17 00:00:00 2001 From: Wojciech Jarmosz Date: Sun, 23 May 2021 20:12:07 +0200 Subject: [PATCH] Add registry saving --- ml_model.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/ml_model.py b/ml_model.py index d8954cf..4b97623 100644 --- a/ml_model.py +++ b/ml_model.py @@ -6,6 +6,8 @@ import os.path import mlflow import sys +from mlflow.tracking import MlflowClient + from tensorflow import keras from tensorflow.keras import layers from tensorflow.keras.layers.experimental import preprocessing @@ -48,7 +50,12 @@ model.fit(train_X, train_Y, verbose=verbose, epochs=epochs) signature = mlflow.models.signature.infer_signature(train_X.numpy(), model.predict(train_X.numpy())) input_data = test_X + +# Dane do rejestracji modelu w MlFlow +mlflow.set_tracking_uri("http://172.17.0.1:5000") +client = MlflowClient() +model_name = "s434704" + with mlflow.start_run(): + mlflow.keras.log_model(model, "movies_on_streaming_platforms_model", registered_model_name="s434704", input_example=input_data.numpy(), signature=signature) mlflow.keras.save_model(model, "movies_on_streaming_platforms_model", input_example=input_data.numpy(), signature=signature) - -