Add registry saving
All checks were successful
s434704-training/pipeline/head This commit looks good
s434704-evaluation/pipeline/head This commit looks good

This commit is contained in:
Wojciech Jarmosz 2021-05-23 20:12:07 +02:00
parent 091c794e5e
commit 8c94666b25

View File

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