mlflow--mlflow
wehub-resource-sync
223 行
7.1 KiB
Python
223 行
7.1 KiB
Python
import json
|
|
import math
|
|
|
|
import keras
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.models import Model
|
|
from mlflow.tracking.fluent import flush_async_logging
|
|
from mlflow.types import Schema, TensorSpec
|
|
from mlflow.utils.autologging_utils import AUTOLOGGING_INTEGRATIONS
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_autologging_config():
|
|
yield
|
|
AUTOLOGGING_INTEGRATIONS.pop("keras", None)
|
|
|
|
|
|
def _create_keras_model():
|
|
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()],
|
|
)
|
|
return model
|
|
|
|
|
|
def test_default_autolog_behavior():
|
|
mlflow.keras.autolog()
|
|
|
|
# Prepare data for a 2-class classification.
|
|
data = np.random.uniform(size=(20, 28, 28, 3))
|
|
label = np.random.randint(2, size=20)
|
|
|
|
model = _create_keras_model()
|
|
|
|
num_epochs = 2
|
|
batch_size = 4
|
|
with mlflow.start_run() as run:
|
|
model.fit(
|
|
data,
|
|
label,
|
|
validation_data=(data, label),
|
|
batch_size=batch_size,
|
|
epochs=num_epochs,
|
|
)
|
|
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 training configs are logged correctly.
|
|
assert int(model_info["batch_size"]) == batch_size
|
|
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 "sparse_categorical_accuracy" in run_metrics
|
|
assert "validation_loss" in run_metrics
|
|
|
|
# Assert metrics are logged in the correct number of times.
|
|
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
|
|
|
|
# Test the loaded pyfunc model produces the same output for the same input as the model.
|
|
test_input = np.random.uniform(size=[2, 28, 28, 3]).astype(np.float32)
|
|
model_uri = f"runs:/{run.info.run_id}/model"
|
|
loaded_pyfunc_model = mlflow.pyfunc.load_model(model_uri)
|
|
np.testing.assert_allclose(
|
|
keras.ops.convert_to_numpy(model(test_input)),
|
|
loaded_pyfunc_model.predict(test_input),
|
|
)
|
|
|
|
# Test the signature is logged.
|
|
input_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 28, 28, 3))])
|
|
output_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 2))])
|
|
mlflow_model = Model.load(model_uri)
|
|
assert mlflow_model.signature.inputs == input_schema
|
|
assert mlflow_model.signature.outputs == output_schema
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"log_every_epoch",
|
|
"log_every_n_steps",
|
|
"log_models",
|
|
"log_model_signatures",
|
|
"save_exported_model",
|
|
),
|
|
[
|
|
(False, 1, False, False, False),
|
|
(False, 2, True, True, True),
|
|
(True, None, False, False, False),
|
|
],
|
|
)
|
|
def test_custom_autolog_behavior(
|
|
log_every_epoch,
|
|
log_every_n_steps,
|
|
log_models,
|
|
log_model_signatures,
|
|
save_exported_model,
|
|
):
|
|
if keras.backend.backend() != "tensorflow" and save_exported_model:
|
|
pytest.skip("Only TensorFlow backend supports saving exported models.")
|
|
mlflow.keras.autolog(
|
|
log_every_epoch=log_every_epoch,
|
|
log_every_n_steps=log_every_n_steps,
|
|
log_models=log_models,
|
|
log_model_signatures=log_model_signatures,
|
|
save_exported_model=save_exported_model,
|
|
)
|
|
|
|
# Prepare data for a 2-class classification.
|
|
data = np.random.uniform(size=(20, 28, 28, 3))
|
|
label = np.random.randint(2, size=20)
|
|
|
|
model = _create_keras_model()
|
|
|
|
num_epochs = 1
|
|
batch_size = 4
|
|
with mlflow.start_run() as run:
|
|
model.fit(
|
|
data,
|
|
label,
|
|
validation_data=(data, label),
|
|
batch_size=batch_size,
|
|
epochs=num_epochs,
|
|
)
|
|
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 training configs are logged correctly.
|
|
assert int(model_info["batch_size"]) == batch_size
|
|
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 "sparse_categorical_accuracy" in run_metrics
|
|
assert "validation_loss" in run_metrics
|
|
|
|
# Assert metrics are logged in the correct number of times.
|
|
loss_history = client.get_metric_history(run_id=run.info.run_id, key="loss")
|
|
if log_every_n_steps:
|
|
metric_length = model.optimizer.iterations.numpy() // log_every_n_steps
|
|
else:
|
|
metric_length = num_epochs
|
|
assert len(loss_history) == metric_length
|
|
|
|
validation_loss_history = client.get_metric_history(
|
|
run_id=run.info.run_id,
|
|
key="validation_loss",
|
|
)
|
|
assert len(validation_loss_history) == num_epochs
|
|
|
|
logged_model = mlflow.last_logged_model()
|
|
if log_models:
|
|
assert logged_model is not None
|
|
assert run_metrics.items() <= {m.key: m.value for m in logged_model.metrics}.items()
|
|
else:
|
|
assert logged_model is None
|
|
assert "mlflow.log-model.history" not in mlflow_run.data.tags
|
|
|
|
|
|
@pytest.mark.parametrize("log_models", [True, False])
|
|
@pytest.mark.parametrize("log_datasets", [True, False])
|
|
def test_keras_autolog_log_datasets(log_datasets, log_models):
|
|
mlflow.keras.autolog(log_datasets=log_datasets, log_models=log_models)
|
|
|
|
# Prepare data for a 2-class classification.
|
|
data = np.random.uniform(size=(20, 28, 28, 3)).astype(np.float32)
|
|
label = np.random.randint(2, size=20)
|
|
|
|
model = _create_keras_model()
|
|
|
|
model.fit(data, label, epochs=2)
|
|
flush_async_logging()
|
|
client = mlflow.MlflowClient()
|
|
run_inputs = client.get_run(mlflow.last_active_run().info.run_id).inputs
|
|
dataset_inputs = run_inputs.dataset_inputs
|
|
if log_datasets:
|
|
assert len(dataset_inputs) == 1
|
|
feature_schema = Schema([
|
|
TensorSpec(np.dtype(np.float32), (-1, 28, 28, 3)),
|
|
])
|
|
target_schema = Schema([
|
|
TensorSpec(np.dtype(np.int64), (-1,)),
|
|
])
|
|
expected = json.dumps({
|
|
"mlflow_tensorspec": {
|
|
"features": feature_schema.to_json(),
|
|
"targets": target_schema.to_json(),
|
|
}
|
|
})
|
|
assert dataset_inputs[0].dataset.schema == expected
|
|
else:
|
|
assert len(dataset_inputs) == 0
|
|
if log_models:
|
|
if log_datasets:
|
|
assert len(run_inputs.model_inputs) == 1
|
|
assert run_inputs.model_inputs[0].model_id == mlflow.last_logged_model().model_id
|
|
else:
|
|
assert mlflow.last_logged_model() is not None
|
|
else:
|
|
assert len(run_inputs.model_inputs) == 0
|