chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user