From a49e24222e4fb0a6d27005861c37b81ace5d5952 Mon Sep 17 00:00:00 2001
From: Wojciech Lidwin <59863785+Halal37@users.noreply.github.com>
Date: Thu, 11 May 2023 20:13:03 +0200
Subject: [PATCH] Add parameters

---
 Dockerfile        |  1 +
 Jenkinsfile_train | 17 ++++++++++++++++-
 ium_train.py      | 15 ++++++++++++---
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 2d3a3d9..26c65cb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -13,6 +13,7 @@ RUN pip3 install numpy
 RUN pip3 install keras
 RUN pip3 install tensorflow
 RUN pip3 install scikit-learn
+RUN pip3 install argparse
 
 
 WORKDIR /app
diff --git a/Jenkinsfile_train b/Jenkinsfile_train
index 55894ab..414f6a8 100644
--- a/Jenkinsfile_train
+++ b/Jenkinsfile_train
@@ -2,6 +2,21 @@ node {
  stage('Preparation') { 
      properties([
          parameters([
+		 	    string(
+                    defaultValue: '20',
+                    description: 'Number of epochs',
+                    name: 'EPOCHS'
+                ),
+                string(
+                    defaultValue: '0.01',
+                    description: 'Learning rate',
+                    name: 'LR'
+                ),
+                string(
+                    defaultValue: '0.2',
+                    description: 'Validation_split',
+                    name: 'VALIDATION_SPLIT'
+                ),
          ])
      ])
  }
@@ -13,7 +28,7 @@ node {
     def testImage = docker.image('s487197/ium:34')
 	testImage.inside{ 
          copyArtifacts filter: 'baltimore_train.csv', projectName: 's487197-create-dataset'
-		 sh "python3 ium_train.py"
+		 sh "python3 ium_train.py" -epochs $EPOCHS -lr $LR -validation_split $VALIDATION_SPLIT"
 	     archiveArtifacts artifacts: 'baltimore_model.h5'
      
 	 
diff --git a/ium_train.py b/ium_train.py
index a88966d..9533bd1 100644
--- a/ium_train.py
+++ b/ium_train.py
@@ -9,6 +9,7 @@ import pandas as pd
 import tensorflow as tf
 import numpy as np
 from sklearn.preprocessing import LabelEncoder
+import argparse
 
 
 def get_x_y(data):
@@ -27,6 +28,14 @@ def get_x_y(data):
 
 
 def train_model():
+
+    parser = argparse.ArgumentParser(description='Train')
+
+    parser.add_argument('-epochs', type=int, default=20)
+    parser.add_argument('-lr', type=float, default=0.01)
+    parser.add_argument('-validation_split', type=int, default=0.2)
+    args = parser.parse_args()
+    
     train = pd.read_csv('baltimore_train.csv')
 
     data_train, x_train, y_train = get_x_y(train)
@@ -38,14 +47,14 @@ def train_model():
     model.add(Dense(10, activation='relu'))
     model.add(Dense(10, activation='relu'))
     model.add(Dense(5, activation="softmax"))
-    model.compile(Adam(learning_rate=0.01), loss='sparse_categorical_crossentropy', metrics = ['accuracy'] )
+    model.compile(Adam(learning_rate=args.lr), loss='sparse_categorical_crossentropy', metrics = ['accuracy'] )
     model.summary()
 
     history = model.fit(
         x_train,
         y_train,
-        epochs=20,
-        validation_split=0.2)
+        epochs=args.epochs,
+        validation_split=args.validation_split)
     hist = pd.DataFrame(history.history)
     hist['epoch'] = history.epoch
     model.save('baltimore_model.h5')