add training param
Some checks failed
s464980-training/pipeline/head There was a failure building this commit

This commit is contained in:
Sheaza 2024-05-14 21:43:04 +02:00
parent 84ef66131c
commit 34ff694fcd
2 changed files with 12 additions and 3 deletions

View File

@ -1,7 +1,7 @@
import pandas as pd import pandas as pd
from tensorflow import keras from tensorflow import keras
from tensorflow.keras import layers from tensorflow.keras import layers
import argparse
class RegressionModel: class RegressionModel:
def __init__(self, optimizer="adam", loss="mean_squared_error"): def __init__(self, optimizer="adam", loss="mean_squared_error"):
@ -42,8 +42,12 @@ class RegressionModel:
self.model.save("model.keras") self.model.save("model.keras")
parser = argparse.ArgumentParser()
parser.add_argument('--epochs')
args = parser.parse_args()
model = RegressionModel() model = RegressionModel()
model.load_data("df_train.csv", "df_test.csv") model.load_data("df_train.csv", "df_test.csv")
model.train() model.train(epochs=args.epochs)
model.evaluate() model.evaluate()
model.save_model() model.save_model()

View File

@ -6,6 +6,11 @@ pipeline {
description: 'Which build to use for copying artifacts', description: 'Which build to use for copying artifacts',
name: 'BUILD_SELECTOR' name: 'BUILD_SELECTOR'
) )
string (
name: 'EPOCHS',
defaultValue: '35',
description: 'Epochs to train'
)
} }
stages { stages {
stage('Checkout repository') { stage('Checkout repository') {
@ -28,7 +33,7 @@ pipeline {
steps { steps {
sh "chmod +x ./train.py" sh "chmod +x ./train.py"
sh "python ./train.py" sh "python ./train.py --epochs ${params.EPOCHS}"
archiveArtifacts artifacts: 'model.keras', onlyIfSuccessful: true archiveArtifacts artifacts: 'model.keras', onlyIfSuccessful: true
} }
} }