mlflow.keras

注意

Autologging 已知与以下软件包版本兼容： 2.13.1 <= tensorflow <= 2.19.0. 当与此范围之外的软件包版本一起使用时，Autologging 可能无法成功。

为 tf.keras 启用自动记录。 请注意，仅支持 tensorflow>=2.3。 例如，试着运行 Keras/TensorFlow example。

对于每个 TensorFlow 模块，autologging 捕获以下信息：

tf.keras

  • 指标 和 参数

  • 训练和验证损失。

  • 用户指定的指标。

  • 优化器配置，例如 learning_rate、momentum 等。

  • 训练配置，例如 epochs、batch_size 等。

  • 工件

  • 训练开始时的模型摘要。

  • 已将 Keras 模型以 MLflow Model 格式保存。

  • 在训练结束时记录 TensorBoard 日志。

tf.keras.callbacks.EarlyStopping

  • 指标 和 参数

  • 来自 EarlyStopping 回调的指标: stopped_epoch, restored_epoch, restore_best_weight, 等

  • fit() 或 fit_generator() 与 EarlyStopping 相关的参数： min_delta, patience, baseline, restore_best_weights, 等

请参阅 autologging 跟踪文档，了解有关 TensorFlow workflows 的更多信息。

注意，autologging 不能与显式的 MLflow 回调一起使用，即 mlflow.tensorflow.MlflowCallback，因为这会导致相同的指标被记录两次。如果你想在回调列表中包含 mlflow.tensorflow.MlflowCallback，请通过调用 mlflow.tensorflow.autolog(disable=True) 关闭 autologging。

param log_models

如果 True，已训练的模型会作为 MLflow 模型工件记录。 如果 False，已训练的模型不会被记录。

param log_datasets

如果 True，数据集信息会记录到 MLflow Tracking。 如果 False，数据集信息不会被记录。

param disable

如果 True，禁用 TensorFlow autologging 集成。如果 False， 启用 TensorFlow autologging 集成。

param exclusive

如果 True，自动记录的内容不会被记录到用户创建的 fluent 运行中。 如果 False，自动记录的内容会被记录到活动的 fluent 运行，后者可能是用户创建的。

param disable_for_unsupported_versions

如果 True，禁用对那些尚未针对此版本的 MLflow 客户端进行测试或与之不兼容的 tensorflow 版本的 autologging。

param silent

如果 True，在 TensorFlow autologging 期间抑制来自 MLflow 的所有事件日志和警告。如果 False，在 TensorFlow autologging 期间显示所有事件和警告。

param registered_model_name

如果提供了该名称，每次训练模型时，都会将其注册为该注册模型的新版本。如果该注册模型不存在，则会创建它。

param log_input_examples

如果 True，在训练期间会收集来自训练数据集的输入示例，并将其与 tf/keras 模型工件一起记录。如果 False，则不会记录输入示例。

param log_model_signatures

如果 True，ModelSignatures 描述模型输入和输出，会在训练期间与 tf/keras 模型工件一起收集并记录。如果 False，则不会记录签名。注意，当将 Pandas DataFrames 传递给 predict() 时，记录带有签名的 TensorFlow 模型会改变它们的 pyfunc 推理行为。当存在签名时，会返回一个 np.ndarray（用于单输出模型）或一个从 str -> np.ndarray 的映射（用于多输出模型）；当不存在签名时，会返回一个 Pandas DataFrame。

param saved_model_kwargs

传递给 tensorflow.saved_model.save 方法的 kwargs 字典。

param keras_model_kwargs

一个要传递给 keras_model.save 方法的 kwargs 字典。

param extra_tags

一个字典，包含要在由 autologging 创建的每个托管运行上设置的额外标签。

param log_every_epoch

如果为 True，则在每个 epoch 结束时记录训练指标。

param log_every_n_steps

如果设置，训练指标将每隔 n 个训练步骤记录一次。 当 log_every_epoch=True 时，log_every_n_steps 必须为 None。

param checkpoint

启用自动模型检查点功能。

param checkpoint_monitor

在自动模型检查点保存中，如果将 model_checkpoint_save_best_only 设置为 True，需要监控的指标名称。

param checkpoint_mode

可选值为 {“min”, “max”} 中的一个。在自动模型检查点保存中，如果 save_best_only=True，是否覆盖当前保存文件的决策是基于对被监控量的最大化或最小化。

param checkpoint_save_best_only

如果为 True，自动模型检查点仅在根据所监控的量认为模型是“最佳”模型时才保存，并且会覆盖之前的检查点模型。

param checkpoint_save_weights_only

在自动模型检查点保存中，如果 True，那么 只会保存模型的权重。否则，优化器状态、 lr-scheduler 状态等也会被加入到检查点中。

param checkpoint_save_freq

“epoch” 或整数。当使用 “epoch” 时，回调在每个 epoch 之后保存模型。当使用整数时，回调会在指定数量的批次结束时保存模型。请注意，如果保存未与 epochs 对齐，监控的指标可能不太可靠（它可能仅反映最少 1 个批次，因为指标在每个 epoch 会被重置）。默认为 “epoch”。

Keras 3 回调，用于将信息记录到 MLflow。

class mlflow.keras.callback.MlflowCallback(log_every_epoch=True, log_every_n_steps=None, model_id=None)[source]

基类： keras.callbacks.Callback

用于将 Keras 的指标/参数/模型等记录到 MLflow 的回调。

此回调在训练开始时记录模型元数据，并在每个 epoch 或每隔用户定义的 n 个步骤记录训练指标到 MLflow。

Parameters

  • log_every_epoch – bool，默认为 True。如果为 True，则在每个 epoch 记录指标。如果为 False，则每 n 步记录指标。

  • log_every_n_steps – int，默认值为 None。如果设置，则每 n 步记录指标。如果为 None，则每个 epoch 记录指标。必须为 None，如果 log_every_epoch=True。

Example 
import keras
import mlflow
import numpy as np

# Prepare data for a 2-class classification.
data = np.random.uniform([8, 28, 28, 3])
label = np.random.randint(2, size=8)
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:
    model.fit(
        data,
        label,
        batch_size=4,
        epochs=2,
        callbacks=[mlflow.keras.MlflowCallback()],
    )
on_batch_end(batch, logs=None)[source]

在每个批次结束时按用户指定的频率记录指标。

on_epoch_end(epoch, logs=None)[source]

在每个训练轮次结束时记录指标。

on_test_end(logs=None)[source]

在验证结束时记录验证指标。

on_train_begin(logs=None)[source]

在训练开始时记录模型架构和优化器配置。

用于加载使用 MLflow 保存的 Keras 模型的函数。

class mlflow.keras.load.KerasModelWrapper(model, signature, save_exported_model=False)[source]

基类： object

get_model_call_method()[source]
get_raw_model()[source]

返回底层模型。

predict(data, **kwargs)[source]
mlflow.keras.load.load_model(model_uri, dst_path=None, custom_objects=None, load_model_kwargs=None)[source]

从 MLflow 加载 Keras 模型。

此方法从 mlflow 加载一个已保存的 Keras 模型，并返回一个 Keras 模型实例。

Parameters

  • model_uri –

    已保存的 Keras 模型在 MLflow 中的 URI。例如：

    • /Users/me/path/to/local/model

    • relative/path/to/local/model

    • s3://my_bucket/path/to/model

    • runs://run-relative/path/to/model

    • models://

    • models://

    有关受支持的 URI 方案的更多信息，请参见 Referencing Artifacts。

  • dst_path – 要下载模型工件的本地文件系统路径。如果未指定，将创建一个本地输出路径。

  • custom_objects – 在 keras.saving.load_model 中的 custom_objects 参数。

  • load_model_kwargs – 用于 keras.saving.load_model 的额外参数。

Example 
import keras
import mlflow
import numpy as np

model = keras.Sequential(
    [
        keras.Input([28, 28, 3]),
        keras.layers.Flatten(),
        keras.layers.Dense(2),
    ]
)
with mlflow.start_run() as run:
    mlflow.keras.log_model(model)

model_url = f"runs:/{run.info.run_id}/{model_path}"
loaded_model = mlflow.keras.load_model(model_url)

# Test the loaded model produces the same output for the same input as the model.
test_input = np.random.uniform(size=[2, 28, 28, 3])
np.testing.assert_allclose(
    keras.ops.convert_to_numpy(model(test_input)),
    loaded_model.predict(test_input),
)
Returns

一个 Keras 模型实例。

用于将 Keras 模型保存到 MLflow 的函数。

mlflow.keras.save.get_default_conda_env()[source]
Returns

默认的 Conda 环境，适用于通过调用 save_model() 和 log_model() 生成的 MLflow Models。

mlflow.keras.save.get_default_pip_requirements()[source]
Returns

Keras 风格生成的 MLflow Models 的默认 pip 依赖项列表。对 save_model() 和 log_model() 的调用会生成一个至少包含这些依赖项的 pip 环境。

mlflow.keras.save.log_model(model, artifact_path: str | None = None, save_exported_model=False, conda_env=None, signature: mlflow.models.signature.ModelSignature = None, input_example: Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple] = None, registered_model_name=None, await_registration_for=300, pip_requirements=None, extra_pip_requirements=None, save_model_kwargs=None, metadata=None, name: str | None = None, params: dict[str, typing.Any] | None = None, tags: dict[str, typing.Any] | None = None, model_type: str | None = None, step: int = 0, model_id: str | None = None)[source]

将 Keras 模型及其元数据记录到 MLflow。

此方法将 Keras 模型及诸如模型签名和 conda 环境等元数据保存到 MLflow。

Parameters

  • model – 一个 keras.Model 的实例。要保存的 Keras 模型。

  • artifact_path – 已弃用。请改用 name。

  • save_exported_model – 默认值为 False。如果为 True，则以导出模型格式保存 Keras 模型，否则以 .keras 格式保存。更多信息请参阅 Keras doc。

  • conda_env –

    可以是 Conda 环境的字典表示，也可以是指向 conda 环境 yaml 文件的路径。如果提供，描述了该模型应运行的环境。至少，应指定包含在 get_default_conda_env() 中的依赖项。如果 None，则会将一个 conda 环境（其 pip 依赖由 mlflow.models.infer_pip_requirements() 推断）添加到模型中。如果依赖推断失败，则回退使用 get_default_pip_requirements。来自 conda_env 的 pip 依赖会被写入 pip 的 requirements.txt 文件，完整的 conda 环境会被写入 conda.yaml。下面是一个 示例 的 conda 环境字典表示：

    {
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "keras==x.y.z"
            ],
        },
    ],
}

  • signature –

    一个 ModelSignature 类的实例，用于描述模型的输入和输出。如果未指定但提供了 input_example，则会根据提供的输入示例和模型自动推断签名。要在提供输入示例时禁用自动签名推断，请将 signature 设置为 False。要手动推断模型签名，请在具有有效模型输入，例如省略了目标列的训练数据集，以及有效模型输出，例如在训练数据集上生成的模型预测的数据集上调用 infer_signature()，例如：

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # 计算模型预测
signature = infer_signature(train, predictions)

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame，然后使用 Pandas 的 split-oriented 格式序列化为 json，或者转换为一个 numpy array，其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

  • registered_model_name – 默认为 None。如果设置，在 registered_model_name 下创建一个模型版本，如果不存在具有该名称的注册模型，也会创建一个注册模型。

  • await_registration_for – 默认为 mlflow.tracking._model_registry.DEFAULT_AWAIT_MAX_SLEEP_SECONDS。等待模型版本完成创建并处于 READY 状态的秒数。默认情况下，该函数等待五分钟。指定 0 或 None 可跳过等待。

  • pip_requirements – 可以是 pip 依赖字符串的可迭代对象（例如 ["keras", "-r requirements.txt", "-c constraints.txt"]）或本地文件系统上 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，该项描述了运行该模型所需的环境。如果 None，则由 mlflow.models.infer_pip_requirements() 从当前软件环境推断出默认的依赖列表。如果依赖推断失败，则回退使用 get_default_pip_requirements。要求和约束会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分进行存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

  • extra_pip_requirements –

    要么是一个 pip 需求字符串的可迭代对象（例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]），要么是本地文件系统上 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，该参数描述了附加的 pip 依赖，这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定：

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip 依赖。

  • save_model_kwargs – 默认值为 None。一个 dict，用于传递给keras.Model.save 方法的 kwargs。

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

  • name – 模型名称。

  • params – 一个用于与模型一同记录的参数字典。

  • tags – 一个要与模型一起记录的标签字典。

  • model_type – 模型的类型。

  • step – 在该步记录模型输出和指标

  • model_id – 模型的 ID。

Example 
import keras
import mlflow

model = keras.Sequential(
    [
        keras.Input([28, 28, 3]),
        keras.layers.Flatten(),
        keras.layers.Dense(2),
    ]
)
with mlflow.start_run() as run:
    mlflow.keras.log_model(model, name="model")
mlflow.keras.save.save_model(model, path, save_exported_model=False, conda_env=None, mlflow_model=None, signature: mlflow.models.signature.ModelSignature = None, input_example: Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple] = None, pip_requirements=None, extra_pip_requirements=None, save_model_kwargs=None, metadata=None)[source]

保存 Keras 模型及其元数据。

此方法将 Keras 模型连同元数据（例如模型签名和 conda 环境）保存到本地文件系统。此方法在 mlflow.keras.log_model() 中被调用。

Parameters

  • model – 一个 keras.Model 的实例。要保存的 Keras 模型。

  • path – 要将 MLflow 模型保存到的本地路径。

  • save_exported_model – 如果为 True，则以导出模型格式保存 Keras 模型；否则以 .keras 格式保存。更多信息，请参阅 https://keras.io/guides/serialization_and_saving/。

  • conda_env –

    可以是 Conda 环境的字典表示，也可以是指向 conda 环境 yaml 文件的路径。如果提供，此项描述了运行该模型应使用的环境。至少，应指定包含在 get_default_conda_env() 中的依赖项。如果 None，则会向模型添加一个其 pip 依赖由 mlflow.models.infer_pip_requirements() 推断的 conda 环境。如果依赖推断失败，则回退使用 get_default_pip_requirements。来自 conda_env 的 pip 依赖将写入 pip 文件 requirements.txt，完整的 conda 环境将写入 conda.yaml。下面是一个 示例 的 conda 环境字典表示：

    {
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "keras==x.y.z"
            ],
        },
    ],
}

  • mlflow_model – 一个 mlflow.models.Model 的实例，默认值为 None。用于添加 Keras 模型元数据的 MLflow 模型配置。如果为 None，则会创建一个空实例。

  • signature –

    一个 ModelSignature 类的实例，用于描述模型的输入和输出。如果未指定但提供了 input_example，则会根据提供的输入示例和模型自动推断签名。要在提供输入示例时禁用自动签名推断，请将 signature 设置为 False。要手动推断模型签名，请在具有有效模型输入，例如省略了目标列的训练数据集，以及有效模型输出，例如在训练数据集上生成的模型预测的数据集上调用 infer_signature()，例如：

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # 计算模型预测
signature = infer_signature(train, predictions)

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame，然后使用 Pandas 的 split-oriented 格式序列化为 json，或者转换为一个 numpy array，其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

  • pip_requirements – 要么是一个 pip 依赖项字符串的可迭代对象（例如 ["keras", "-r requirements.txt", "-c constraints.txt"]）或本地文件系统上 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，该值描述了应该在其中运行此模型的环境。如果 None，则由 mlflow.models.infer_pip_requirements() 从当前软件环境推断出一个默认的依赖项列表。如果依赖项推断失败，则回退为使用 get_default_pip_requirements。依赖项和约束会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

  • extra_pip_requirements –

    要么是一个 pip 需求字符串的可迭代对象（例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]），要么是本地文件系统上 pip requirements 文件的字符串路径（例如 "requirements.txt"）。如果提供，该参数描述了附加的 pip 依赖，这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定：

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip 依赖。

  • save_model_kwargs – 一个要传递给 keras.Model.save 方法的 kwargs 字典。

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

Example 
import keras
import mlflow

model = keras.Sequential(
    [
        keras.Input([28, 28, 3]),
        keras.layers.Flatten(),
        keras.layers.Dense(2),
    ]
)
with mlflow.start_run() as run:
    mlflow.keras.save_model(model, "./model")