wehub-resource-sync
114 lines
3.6 KiB
Python
114 lines
3.6 KiB
Python
import math
|
|
|
|
import keras
|
|
import numpy as np
|
|
|
|
import mlflow
|
|
from mlflow.keras.callback import MlflowCallback
|
|
from mlflow.tracking.fluent import flush_async_logging
|
|
|
|
|
|
def test_keras_mlflow_callback_log_every_epoch():
|
|
# Prepare data for a 2-class classification.
|
|
data = np.random.uniform(size=(20, 28, 28, 3))
|
|
label = np.random.randint(2, size=20)
|
|
|
|
model = keras.Sequential([
|
|
keras.Input([28, 28, 3]),
|
|
keras.layers.Flatten(),
|
|
keras.layers.Dense(2),
|
|
])
|
|
|
|
model.compile(
|
|
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
|
|
optimizer=keras.optimizers.Adam(0.001),
|
|
metrics=[keras.metrics.SparseCategoricalAccuracy()],
|
|
)
|
|
|
|
num_epochs = 2
|
|
with mlflow.start_run() as run:
|
|
mlflow_callback = MlflowCallback(log_every_epoch=True)
|
|
model.fit(
|
|
data,
|
|
label,
|
|
validation_data=(data, label),
|
|
batch_size=4,
|
|
epochs=num_epochs,
|
|
callbacks=[mlflow_callback],
|
|
)
|
|
flush_async_logging()
|
|
client = mlflow.MlflowClient()
|
|
mlflow_run = client.get_run(run.info.run_id)
|
|
run_metrics = mlflow_run.data.metrics
|
|
model_info = mlflow_run.data.params
|
|
|
|
assert "sparse_categorical_accuracy" in run_metrics
|
|
assert model_info["optimizer_name"] == "adam"
|
|
assert math.isclose(float(model_info["optimizer_learning_rate"]), 0.001, rel_tol=1e-6)
|
|
assert "loss" in run_metrics
|
|
assert "validation_loss" in run_metrics
|
|
|
|
loss_history = client.get_metric_history(run_id=run.info.run_id, key="loss")
|
|
assert len(loss_history) == num_epochs
|
|
|
|
validation_loss_history = client.get_metric_history(
|
|
run_id=run.info.run_id,
|
|
key="validation_loss",
|
|
)
|
|
assert len(validation_loss_history) == num_epochs
|
|
|
|
|
|
def test_keras_mlflow_callback_log_every_n_steps():
|
|
# Prepare data for a 2-class classification.
|
|
data = np.random.uniform(size=(20, 28, 28, 3))
|
|
label = np.random.randint(2, size=20)
|
|
|
|
model = keras.Sequential([
|
|
keras.Input([28, 28, 3]),
|
|
keras.layers.Flatten(),
|
|
keras.layers.Dense(2),
|
|
])
|
|
|
|
model.compile(
|
|
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
|
|
optimizer=keras.optimizers.Adam(0.001),
|
|
metrics=[keras.metrics.SparseCategoricalAccuracy()],
|
|
)
|
|
|
|
log_every_n_steps = 1
|
|
num_epochs = 2
|
|
with mlflow.start_run() as run:
|
|
mlflow_callback = MlflowCallback(log_every_epoch=False, log_every_n_steps=log_every_n_steps)
|
|
model.fit(
|
|
data,
|
|
label,
|
|
validation_data=(data, label),
|
|
batch_size=4,
|
|
epochs=num_epochs,
|
|
callbacks=[mlflow_callback],
|
|
)
|
|
flush_async_logging()
|
|
client = mlflow.MlflowClient()
|
|
mlflow_run = client.get_run(run.info.run_id)
|
|
run_metrics = mlflow_run.data.metrics
|
|
model_info = mlflow_run.data.params
|
|
|
|
assert "sparse_categorical_accuracy" in run_metrics
|
|
assert model_info["optimizer_name"] == "adam"
|
|
assert math.isclose(float(model_info["optimizer_learning_rate"]), 0.001, rel_tol=1e-6)
|
|
assert "loss" in run_metrics
|
|
assert "validation_loss" in run_metrics
|
|
|
|
loss_history = client.get_metric_history(run_id=run.info.run_id, key="loss")
|
|
assert len(loss_history) == model.optimizer.iterations.numpy() // log_every_n_steps
|
|
|
|
validation_loss_history = client.get_metric_history(
|
|
run_id=run.info.run_id,
|
|
key="validation_loss",
|
|
)
|
|
assert len(validation_loss_history) == num_epochs
|
|
|
|
|
|
def test_old_callback_still_exists():
|
|
assert mlflow.keras.MLflowCallback is mlflow.keras.MlflowCallback
|