Saving signatures and input example in train_mlflow.py

Jenkinsfile_train

@@ -34,8 +34,8 @@ pipeline {
                         sh 'chmod +x dlgssdpytorch.py'
                         sh 'python3 ./dlgssdpytorch.py $PARAMETRY'
                         sh 'chmod +x train_mlflow.py'
-						sh 'chmod +x generate_MLmodel.py'
+						//sh 'chmod +x generate_MLmodel.py'
-						sh 'python3 ./generate_MLmodel.py'
+						//sh 'python3 ./generate_MLmodel.py'
                         //sh 'mlflow ./train_mlflow.py'
 						//sh 'mlflow run .'
                     }

train_mlflow.py

@@ -8,6 +8,7 @@ import argparse
 import mlflow
 import mlflow.pytorch
 from urllib.parse import urlparse
+from mlflow.models.signature import infer_signature
 
 
 class LayerLinearRegression(nn.Module):
@@ -28,6 +29,7 @@ args = parser.parse_args()
 if __name__ == "__main__":
     lr = args.lr
     n_epochs = args.epochs
+    mlflow.set_experiment("s426206")
     with mlflow.start_run():
         mlflow.log_param("lr", lr)
         mlflow.log_param("epochs", n_epochs)
@@ -119,15 +121,19 @@ if __name__ == "__main__":
                         'loss': lr,
                     }, 'model.pt')
 
-        tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
+        x_train = np.array(train_dataset)[:,0] #(Sales Sum row)
+        input_example = np.reshape(x_train, (-1,1))
 
+        with torch.no_grad():        
+            model.eval()
+            siganture = infer_signature(x_train, model(torch.tensor(np.reshape(x_train, (-1,1))).float()).numpy())
+
+        #mlflow.set_experiment("s426206")
+        #mlflow.set_tracking_uri("http://172.17.0.1:5000")
+        tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
+        # print(tracking_url_type_store)
         # Model registry does not work with file store
         if tracking_url_type_store != "file":
-
+            mlflow.sklearn.log_model(model, "model", registered_model_name="s426206", signature=siganture, input_example=input_example)
-            # Register the model
-            # There are other ways to use the Model Registry, which depends on the use case,
-            # please refer to the doc for more information:
-            # https://mlflow.org/docs/latest/model-registry.html#api-workflow
-            mlflow.sklearn.log_model(model, "model", registered_model_name="ElasticnetWineModel")
         else:
-            mlflow.sklearn.log_model(model, "model")
+            mlflow.sklearn.log_model(model, "model", signature=siganture, input_example=input_example)