ium_470607/lab5/train/train.py
2021-05-16 23:01:34 +02:00

63 lines
1.8 KiB
Python

from datetime import datetime
import mlflow
import pandas as pd
from sacred import Experiment
from sacred.observers import FileStorageObserver, MongoObserver
import sys
import tensorflow
from tensorflow.keras import layers
ex = Experiment("470607", interactive=False, save_git_info=False)
ex.observers.append(
MongoObserver(url='mongodb://mongo_user:mongo_password_IUM_2021@172.17.0.1:27017', db_name='sacred'))
ex.observers.append(FileStorageObserver('my_runs'))
@ex.config
def my_config():
learning_rate = float(sys.argv[1])
@ex.capture
def prepare_train_model(learning_rate, _run):
with mlflow.start_run():
_run.info["prepare_model"] = str(datetime.now())
X_train = pd.read_csv('train.csv')
X_valid = pd.read_csv('valid.csv')
Y_train = X_train.pop('stabf')
Y_train = pd.get_dummies(Y_train)
Y_valid = X_valid.pop('stabf')
Y_valid = pd.get_dummies(Y_valid)
model = tensorflow.keras.Sequential([
layers.Input(shape=(12,)),
layers.Dense(32),
layers.Dense(16),
layers.Dense(2, activation='softmax')
])
model.compile(
loss=tensorflow.keras.losses.BinaryCrossentropy(),
optimizer=tensorflow.keras.optimizers.Adam(learning_rate=learning_rate),
metrics=[tensorflow.keras.metrics.BinaryAccuracy()])
history = model.fit(X_train, Y_train, epochs=2, validation_data=(X_valid, Y_valid))
model.save('grid-stability-dense.h5')
_run.info['history'] = str(history.history['loss'][-1])
mlflow.log_metric('loss', history.history['loss'][-1])
mlflow.log_param('learning_rate', learning_rate)
@ex.main
def my_main(learning_rate):
print(prepare_train_model())
r = ex.run()
ex.add_artifact('grid-stability-dense.h5')