From 858e9ec215b82a19a2cdca5e10daf562476a5b0d Mon Sep 17 00:00:00 2001 From: Marcin Kostrzewski Date: Mon, 9 May 2022 18:23:12 +0200 Subject: [PATCH] Archive MLFlow artifacts --- Jenkinsfile-train | 2 +- train_model.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/Jenkinsfile-train b/Jenkinsfile-train index de37f9d..8995392 100644 --- a/Jenkinsfile-train +++ b/Jenkinsfile-train @@ -41,7 +41,7 @@ pipeline { stage('Archive model and evaluate it') { steps { archiveArtifacts artifacts: 'model_out', onlyIfSuccessful: true - archiveArtifacts artifacts: 'sacred_runs/**', onlyIfSuccessful: true + archiveArtifacts artifacts: 'mlruns/**', onlyIfSuccessful: true build job: "s444409-evaluation/${params.BRANCH}/", parameters: [string(name: 'BRANCH', value: "${params.BRANCH}")] } } diff --git a/train_model.py b/train_model.py index 5bdd938..d3af166 100644 --- a/train_model.py +++ b/train_model.py @@ -1,6 +1,7 @@ import torch import mlflow import argparse +import numpy as np from torch import nn from torch.utils.data import DataLoader from urllib.parse import urlparse @@ -34,6 +35,8 @@ if __name__ == "__main__": plant_test = PlantsDataset('data/Plant_1_Generation_Data.csv.test') plant_train = PlantsDataset('data/Plant_1_Generation_Data.csv.train') + input_example = np.array([plant_test.x_train.numpy()[0]]) + train_dataloader = DataLoader(plant_train, batch_size=batch_size) test_dataloader = DataLoader(plant_test, batch_size=batch_size) @@ -71,9 +74,10 @@ if __name__ == "__main__": if tracking_url_type_store != "file": mlflow.pytorch.log_model( model, - "s444409-power-plant-model", - registered_model_name="s444409PowerPlant", - signature=signature + "s444409", + registered_model_name="s444409", + signature=signature, + input_example=input_example ) else: - mlflow.pytorch.log_model(model, "s444409-power-plant-model", signature=signature) + mlflow.pytorch.log_model(model, "s444409", signature=signature, input_example=input_example)