This commit is contained in:
Jakub Zaręba 2023-05-10 20:42:51 +02:00
parent 8f458fd54a
commit 134226fa79

100
train.py
View File

@ -1,69 +1,69 @@
from sacred import Experiment from sacred import Experiment
from sacred.observers import MongoObserver, FileStorageObserver from sacred.observers import MongoObserver, FileStorageObserver
ex = Experiment('s487187-training') # ex = Experiment('s487187-training')
ex.observers.append(MongoObserver(url='mongodb://admin:IUM_2021@172.17.0.1:27017', db_name='sacred')) # ex.observers.append(MongoObserver(url='mongodb://admin:IUM_2021@172.17.0.1:27017', db_name='sacred'))
ex.observers.append(FileStorageObserver('results')) # ex.observers.append(FileStorageObserver('results'))
ex.use_git = False # ex.use_git = False
@ex.config # @ex.config
def my_config(): # def my_config():
data_file = 'data.csv' data_file = 'data.csv'
model_file = 'model.h5' model_file = 'model.h5'
epochs = 10 epochs = 10
batch_size = 32 batch_size = 32
test_size = 0.2 test_size = 0.2
random_state = 42 random_state = 42
@ex.capture # @ex.capture
def train_model(data_file, model_file, epochs, batch_size, test_size, random_state): # def train_model(data_file, model_file, epochs, batch_size, test_size, random_state):
import pandas as pd import pandas as pd
from sklearn.model_selection import train_test_split from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler from sklearn.preprocessing import MinMaxScaler
import tensorflow as tf import tensorflow as tf
from imblearn.over_sampling import SMOTE from imblearn.over_sampling import SMOTE
smote = SMOTE(random_state=random_state) smote = SMOTE(random_state=random_state)
data = pd.read_csv(data_file, sep=';') data = pd.read_csv(data_file, sep=';')
print('Total rows:', len(data)) print('Total rows:', len(data))
print('Rows with medal:', len(data.dropna(subset=['Medal']))) print('Rows with medal:', len(data.dropna(subset=['Medal'])))
data = pd.get_dummies(data, columns=['Sex', 'Medal']) data = pd.get_dummies(data, columns=['Sex', 'Medal'])
data = data.drop(columns=['Name', 'Team', 'NOC', 'Games', 'Year', 'Season', 'City', 'Sport', 'Event']) data = data.drop(columns=['Name', 'Team', 'NOC', 'Games', 'Year', 'Season', 'City', 'Sport', 'Event'])
scaler = MinMaxScaler() scaler = MinMaxScaler()
data = pd.DataFrame(scaler.fit_transform(data), columns=data.columns) data = pd.DataFrame(scaler.fit_transform(data), columns=data.columns)
X = data.filter(regex='Sex|Age') X = data.filter(regex='Sex|Age')
y = data.filter(regex='Medal') y = data.filter(regex='Medal')
y = pd.get_dummies(y) y = pd.get_dummies(y)
X = X.fillna(0) X = X.fillna(0)
y = y.fillna(0) y = y.fillna(0)
y = y.values y = y.values
X_resampled, y_resampled = smote.fit_resample(X, y) X_resampled, y_resampled = smote.fit_resample(X, y)
X_train, X_test, y_train, y_test = train_test_split(X_resampled, y_resampled, test_size=test_size, random_state=random_state) X_train, X_test, y_train, y_test = train_test_split(X_resampled, y_resampled, test_size=test_size, random_state=random_state)
model = tf.keras.models.Sequential() model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Dense(64, input_dim=X_train.shape[1], activation='relu')) model.add(tf.keras.layers.Dense(64, input_dim=X_train.shape[1], activation='relu'))
model.add(tf.keras.layers.Dense(32, activation='relu')) model.add(tf.keras.layers.Dense(32, activation='relu'))
model.add(tf.keras.layers.Dense(y.shape[1], activation='softmax')) model.add(tf.keras.layers.Dense(y.shape[1], activation='softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy']) model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size) model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size)
loss, accuracy = model.evaluate(X_test, y_test) loss, accuracy = model.evaluate(X_test, y_test)
print('Test accuracy:', accuracy) print('Test accuracy:', accuracy)
model.save(model_file) model.save(model_file)
return accuracy # return accuracy
@ex.main # @ex.main
def run_experiment(): # def run_experiment():
accuracy = train_model() # accuracy = train_model()
ex.log_scalar('accuracy', accuracy) # ex.log_scalar('accuracy', accuracy)
ex.add_artifact('model.h5') # ex.add_artifact('model.h5')