wehub-resource-sync
58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
import numpy as np
|
|
import pytest
|
|
import tensorflow as tf
|
|
from tensorflow import keras
|
|
|
|
import mlflow
|
|
from mlflow.tensorflow.callback import MlflowCallback
|
|
|
|
|
|
@pytest.mark.parametrize(("log_every_epoch", "log_every_n_steps"), [(True, None), (False, 1)])
|
|
def test_tf_mlflow_callback(log_every_epoch, log_every_n_steps):
|
|
# Prepare data for a 2-class classification.
|
|
data = tf.random.uniform([20, 28, 28, 3])
|
|
label = tf.convert_to_tensor(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()],
|
|
)
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow_callback = MlflowCallback(
|
|
run=run,
|
|
log_every_epoch=log_every_epoch,
|
|
log_every_n_steps=log_every_n_steps,
|
|
)
|
|
model.fit(
|
|
data,
|
|
label,
|
|
validation_data=(data, label),
|
|
batch_size=4,
|
|
# Increase the epochs size so that logs
|
|
# are flushed correctly
|
|
epochs=5,
|
|
callbacks=[mlflow_callback],
|
|
)
|
|
|
|
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 "loss" in run_metrics
|
|
assert "sparse_categorical_accuracy" in run_metrics
|
|
assert model_info["optimizer_name"].lower() == "adam"
|
|
np.testing.assert_almost_equal(float(model_info["optimizer_learning_rate"]), 0.001)
|
|
|
|
|
|
def test_old_callback_still_exists():
|
|
assert mlflow.tensorflow.MLflowCallback is mlflow.tensorflow.MlflowCallback
|