Add parameters

This commit is contained in:
Wojciech Lidwin 2023-05-11 20:13:03 +02:00
parent 023f843d90
commit a49e24222e
3 changed files with 29 additions and 4 deletions

View File

@ -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

View File

@ -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'

View File

@ -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')