From 6b4ac85f8af09ba58fefbdc283b386accb212627 Mon Sep 17 00:00:00 2001 From: Adrian Charkiewicz Date: Sun, 15 May 2022 20:06:19 +0200 Subject: [PATCH] signature --- pytorch/pytorch.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/pytorch/pytorch.py b/pytorch/pytorch.py index 34b8c94..6c537d2 100644 --- a/pytorch/pytorch.py +++ b/pytorch/pytorch.py @@ -22,7 +22,8 @@ from sacred import Experiment from sacred.observers import FileStorageObserver from sacred.observers import MongoObserver import mlflow -import mlflow.keras +import mlflow.pytorch +from mlflow.models import infer_signature # In[2]: ex = Experiment(save_git_info=False) @@ -185,16 +186,17 @@ def fit(epochs, lr, model, train_loader, val_loader, _log, _run, opt_func=torch. _run.info["epochs"] = epochs - signature = mlflow.models.signature.infer_signature(house_price_features, linear_model.predict(house_price_features)) + signature = mlflow.models.signature.infer_signature(train_ds) tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme - sampleInp = [0.1, 0.1, 546.0, 555.223, 1., 1., 33.16376, 84.12426] + input_example = [0.1, 0.1, 546.0, 555.223, 1., 1., 33.16376, 84.12426] if tracking_url_type_store != "file": - mlflow.keras.log_model(model, "model", registered_model_name="red-wine-quality", signature=signature) + mlflow.pytorch.log_model(model, "model", registered_model_name="s444354", signature=siganture, input_example=input_example) else: - mlflow.keras.log_model(model, "model", signature=signature, input_example=np.array(sampleInp)) + mlflow.pytorch.log_model(model, "model", signature=siganture, input_example=input_example) + mlflow.pytorch.save_model(model, "my_model", signature=siganture, input_example=input_example)