wehub-resource-sync
170 lines
6.6 KiB
Python
170 lines
6.6 KiB
Python
"""
|
|
Train a simple Keras DL model on the dataset used in MLflow tutorial (wine-quality.csv).
|
|
|
|
Dataset is split into train (~ 0.56), validation(~ 0.19) and test (0.25).
|
|
Validation data is used to select the best hyperparameters, test set performance is evaluated only
|
|
at epochs which improved performance on the validation dataset. The model with best validation set
|
|
performance is logged with MLflow.
|
|
"""
|
|
|
|
import math
|
|
import warnings
|
|
|
|
import click
|
|
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.metrics import mean_squared_error
|
|
from sklearn.model_selection import train_test_split
|
|
from tensorflow import keras
|
|
from tensorflow.keras.callbacks import Callback
|
|
from tensorflow.keras.layers import Dense, Lambda
|
|
from tensorflow.keras.models import Sequential
|
|
from tensorflow.keras.optimizers import SGD
|
|
|
|
import mlflow
|
|
from mlflow.models import infer_signature
|
|
|
|
|
|
def eval_and_log_metrics(prefix, actual, pred, epoch):
|
|
rmse = np.sqrt(mean_squared_error(actual, pred))
|
|
mlflow.log_metric(f"{prefix}_rmse", rmse, step=epoch)
|
|
return rmse
|
|
|
|
|
|
def get_standardize_f(train):
|
|
mu = np.mean(train, axis=0)
|
|
std = np.std(train, axis=0)
|
|
return lambda x: (x - mu) / std
|
|
|
|
|
|
class MlflowCheckpoint(Callback):
|
|
"""
|
|
Example of Keras MLflow logger.
|
|
Logs training metrics and final model with MLflow.
|
|
|
|
We log metrics provided by Keras during training and keep track of the best model (best loss
|
|
on validation dataset). Every improvement of the best model is also evaluated on the test set.
|
|
|
|
At the end of the training, log the best model with MLflow.
|
|
"""
|
|
|
|
def __init__(self, test_x, test_y, loss="rmse"):
|
|
self._test_x = test_x
|
|
self._test_y = test_y
|
|
self.train_loss = f"train_{loss}"
|
|
self.val_loss = f"val_{loss}"
|
|
self.test_loss = f"test_{loss}"
|
|
self._best_train_loss = math.inf
|
|
self._best_val_loss = math.inf
|
|
self._best_model = None
|
|
self._next_step = 0
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""
|
|
Log the best model at the end of the training run.
|
|
"""
|
|
if not self._best_model:
|
|
raise Exception("Failed to build any model")
|
|
mlflow.log_metric(self.train_loss, self._best_train_loss, step=self._next_step)
|
|
mlflow.log_metric(self.val_loss, self._best_val_loss, step=self._next_step)
|
|
predictions = self._best_model.predict(self._test_x)
|
|
signature = infer_signature(self._test_x, predictions)
|
|
mlflow.tensorflow.log_model(self._best_model, name="model", signature=signature)
|
|
|
|
def on_epoch_end(self, epoch, logs=None):
|
|
"""
|
|
Log Keras metrics with MLflow. If model improved on the validation data, evaluate it on
|
|
a test set and store it as the best model.
|
|
"""
|
|
if not logs:
|
|
return
|
|
self._next_step = epoch + 1
|
|
train_loss = logs["loss"]
|
|
val_loss = logs["val_loss"]
|
|
mlflow.log_metrics({self.train_loss: train_loss, self.val_loss: val_loss}, step=epoch)
|
|
|
|
if val_loss < self._best_val_loss:
|
|
# The result improved in the validation set.
|
|
# Log the model with mlflow and also evaluate and log on test set.
|
|
self._best_train_loss = train_loss
|
|
self._best_val_loss = val_loss
|
|
self._best_model = keras.models.clone_model(self.model)
|
|
self._best_model.set_weights([x.copy() for x in self.model.get_weights()])
|
|
preds = self._best_model.predict(self._test_x)
|
|
eval_and_log_metrics("test", self._test_y, preds, epoch)
|
|
|
|
|
|
@click.command(
|
|
help="Trains an Keras model on wine-quality dataset. "
|
|
"The input is expected in csv format. "
|
|
"The model and its metrics are logged with mlflow."
|
|
)
|
|
@click.option("--epochs", type=click.INT, default=100, help="Maximum number of epochs to evaluate.")
|
|
@click.option(
|
|
"--batch-size", type=click.INT, default=16, help="Batch size passed to the learning algo."
|
|
)
|
|
@click.option("--learning-rate", type=click.FLOAT, default=1e-2, help="Learning rate.")
|
|
@click.option("--momentum", type=click.FLOAT, default=0.9, help="SGD momentum.")
|
|
@click.option("--seed", type=click.INT, default=97531, help="Seed for the random generator.")
|
|
@click.argument("training_data")
|
|
def run(training_data, epochs, batch_size, learning_rate, momentum, seed):
|
|
warnings.filterwarnings("ignore")
|
|
data = pd.read_csv(training_data, sep=";")
|
|
# Split the data into training and test sets. (0.75, 0.25) split.
|
|
train, test = train_test_split(data, random_state=seed)
|
|
train, valid = train_test_split(train, random_state=seed)
|
|
# The predicted column is "quality" which is a scalar from [3, 9]
|
|
train_x = train.drop(["quality"], axis=1).astype("float32").values
|
|
train_y = train[["quality"]].astype("float32").values
|
|
valid_x = valid.drop(["quality"], axis=1).astype("float32").values
|
|
|
|
valid_y = valid[["quality"]].astype("float32").values
|
|
|
|
test_x = test.drop(["quality"], axis=1).astype("float32").values
|
|
test_y = test[["quality"]].astype("float32").values
|
|
|
|
with mlflow.start_run():
|
|
if epochs == 0: # score null model
|
|
eval_and_log_metrics(
|
|
"train", train_y, np.ones(len(train_y)) * np.mean(train_y), epoch=-1
|
|
)
|
|
eval_and_log_metrics("val", valid_y, np.ones(len(valid_y)) * np.mean(valid_y), epoch=-1)
|
|
eval_and_log_metrics("test", test_y, np.ones(len(test_y)) * np.mean(test_y), epoch=-1)
|
|
else:
|
|
with MlflowCheckpoint(test_x, test_y) as mlflow_logger:
|
|
model = Sequential()
|
|
model.add(Lambda(get_standardize_f(train_x)))
|
|
model.add(
|
|
Dense(
|
|
train_x.shape[1],
|
|
activation="relu",
|
|
kernel_initializer="normal",
|
|
input_shape=(train_x.shape[1],),
|
|
)
|
|
)
|
|
model.add(Dense(16, activation="relu", kernel_initializer="normal"))
|
|
model.add(Dense(16, activation="relu", kernel_initializer="normal"))
|
|
model.add(Dense(1, kernel_initializer="normal", activation="linear"))
|
|
model.compile(
|
|
loss="mean_squared_error",
|
|
optimizer=SGD(lr=learning_rate, momentum=momentum),
|
|
metrics=[],
|
|
)
|
|
|
|
model.fit(
|
|
train_x,
|
|
train_y,
|
|
batch_size=batch_size,
|
|
epochs=epochs,
|
|
verbose=1,
|
|
validation_data=(valid_x, valid_y),
|
|
callbacks=[mlflow_logger],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run()
|