add epochs to py
This commit is contained in:
parent
05291f41f6
commit
e24b6eef5a
|
@ -1,5 +1,5 @@
|
|||
matrix:
|
||||
MSTEPS:
|
||||
EPOCHS:
|
||||
- 20
|
||||
- 30
|
||||
|
||||
|
|
6
run.py
6
run.py
|
@ -8,6 +8,10 @@ import os
|
|||
from keras import backend as K
|
||||
K.clear_session()
|
||||
|
||||
EPOCHS = int(os.getenv("EPOCHS", default = 10))
|
||||
|
||||
|
||||
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
# if gpus:
|
||||
# try:
|
||||
|
@ -71,7 +75,7 @@ with tf.device('/GPU:0'):
|
|||
model.compile(optimizer='SGD',
|
||||
loss='categorical_crossentropy',
|
||||
metrics=['accuracy'])
|
||||
model.fit(X_train_scaled, y_train_categorical, epochs=25)
|
||||
model.fit(X_train_scaled, y_train_categorical, epochs=EPOCHS)
|
||||
model.save('mymodel.keras')
|
||||
|
||||
print("finished training")
|
||||
|
|
Loading…
Reference in New Issue