diff --git a/train.py b/train.py index f23ee93..8cbbb95 100644 --- a/train.py +++ b/train.py @@ -1,7 +1,7 @@ import pandas as pd from tensorflow import keras from tensorflow.keras import layers - +import argparse class RegressionModel: def __init__(self, optimizer="adam", loss="mean_squared_error"): @@ -42,8 +42,12 @@ class RegressionModel: self.model.save("model.keras") +parser = argparse.ArgumentParser() +parser.add_argument('--epochs') + +args = parser.parse_args() model = RegressionModel() model.load_data("df_train.csv", "df_test.csv") -model.train() +model.train(epochs=args.epochs) model.evaluate() model.save_model() diff --git a/training/Jenkinsfile b/training/Jenkinsfile index 8f118a7..3950008 100644 --- a/training/Jenkinsfile +++ b/training/Jenkinsfile @@ -6,6 +6,11 @@ pipeline { description: 'Which build to use for copying artifacts', name: 'BUILD_SELECTOR' ) + string ( + name: 'EPOCHS', + defaultValue: '35', + description: 'Epochs to train' + ) } stages { stage('Checkout repository') { @@ -28,7 +33,7 @@ pipeline { steps { sh "chmod +x ./train.py" - sh "python ./train.py" + sh "python ./train.py --epochs ${params.EPOCHS}" archiveArtifacts artifacts: 'model.keras', onlyIfSuccessful: true } }