mlflow.keras

注意

已知 autologging 与以下包版本兼容： 3.0.2 <= keras <= 3.12.0。使用此范围外的包版本时，autologging 可能不会成功。

为 Keras 启用 autologging。

此方法配置 Keras 工作流的 autologging。仅支持 Keras > 3。对于较低版本的 Keras（也称为 tf-keras）的使用，请参阅 mlflow.tensorflow flavor。总的来说，调用此 mlflow.keras.autolog() 函数会将 keras.Model.fit 方法替换为 MLflow 提供的自定义 fit 方法，该方法会在相应的时间将指标/参数/信息/模型记录到 MLflow。

Autologging 与 Keras 支持的所有后端兼容，包括 Tensorflow、PyTorch 和 JAX。

请注意，autologging 仅在您使用 model.fit() 进行训练时有效。如果您正在编写自定义训练循环，则需要使用手动日志记录。

param log_every_epoch: 如果为 True，则将在每个 epoch 结束时记录训练指标。
param log_every_n_steps: 如果设置，则每 n 个训练步骤记录一次训练指标。log_every_n_steps 在 log_every_epoch=True 时必须为 None。
param log_models: 如果为 True，则将在 model.fit() 结束时将 Keras 模型记录到 MLflow。
param log_model_signatures: 如果为 True，将自动捕获并记录模型签名。
param save_exported_model: 如果为 True，模型将以导出的格式（编译后的图）保存，适用于服务和部署。如果为 False，模型将以 .keras 格式保存，该格式包含模型架构和权重。
param log_datasets: 如果为 True，将把数据集元数据记录到 MLflow。
param log_input_examples: 如果为 True，将记录输入示例。
param disable: 如果为 True，则禁用 Keras autologging。
param exclusive: 如果为 True，则 autologged 内容不会记录到用户创建的 fluent runs。如果为 False，则 autologged 内容会记录到活动的 fluent run，该 run 可能是用户创建的。disable_for_unsupported_versions：如果为 True，则禁用与不支持的版本 autologging。
param disable_for_unsupported_versions: 如果为 True，则禁用与此 MLflow 客户端版本未经测试或不兼容的 keras 版本 autologging。
param silent: 如果为 True，则在 Keras autologging 期间抑制 MLflow 的所有事件日志和警告。如果为 True，则在 Keras autologging 期间显示所有事件和警告。
param registered_model_name: 如果设置，则每次训练模型时，它都会作为注册模型的新模型版本进行注册。如果注册模型不存在，则会创建该模型。
param save_model_kwargs: 传递给 keras.Model.save() 的额外 kwargs。
param extra_tags: 要在 autologging 创建的每个托管 run 上设置的额外标签字典。

示例

import keras
import mlflow
import numpy as np

mlflow.keras.autolog()

# 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)

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

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

Bases: keras.src.callbacks.callback.Callback

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

此回调会在训练开始时记录模型元数据，并根据用户定义的频率（每个 epoch 或每 n 步）将训练指标记录到 MLflow。

参数

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

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]: 根据用户指定的频率在每个 batch 结束时记录指标。

on_epoch_end(epoch, logs=None)[source]: 在每个 epoch 结束时记录指标。

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]

Bases: 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 模型实例。

参数

model_uri –
MLflow 中保存的 Keras 模型的 URI。例如：
- /Users/me/path/to/local/model
- relative/path/to/local/model
- s3://my_bucket/path/to/model
- runs:/<mlflow_run_id>/run-relative/path/to/model
- models:/<model_name>/<model_version>
- models:/<model_name>/<stage>
For more information about supported URI schemes, see Referencing Artifacts.
dst_path – 要将模型 artifact 下载到的本地文件系统路径。如果未指定，将创建一个本地输出路径。
custom_objects – keras.saving.load_model 中的 custom_objects 参数。
load_model_kwargs – keras.saving.load_model 的额外参数。

示例

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),
)

返回: 一个 Keras 模型实例。

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

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

返回: 通过 save_model() 和 log_model() 调用生成的 MLflow Models 的默认 Conda 环境。

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

返回: Keras flavor 生成的 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。

参数

model – keras.Model 的一个实例。要保存的 Keras 模型。
artifact_path – Deprecated. Use name instead.
save_exported_model – 默认为 False。如果为 True，则以导出的模型格式保存 Keras 模型，否则以 .keras 格式保存。有关更多信息，请参阅 Keras 文档。
conda_env –
Conda 环境的字典表示形式或 conda 环境 yaml 文件的路径。如果提供，它将描述此模型应运行的环境。至少，它应指定 get_default_conda_env() 中包含的依赖项。如果为 None，则将一个通过 mlflow.models.infer_pip_requirements() 推断出的 pip 需求的 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"
            ],
        },
    ],
}
```
signature –
一个 ModelSignature 类的实例，描述了模型的输入和输出。如果未指定但提供了 input_example，将根据提供的输入示例和模型自动推断签名。要在使用输入示例时禁用自动签名推断，请将 signature 设置为 False。要手动推断模型签名，请在具有有效模型输入的模型（例如，省略目标列的训练数据集）和有效模型输出（例如，在训练数据集上进行的模型预测）的数据集上调用 infer_signature()，例如：
```
from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # compute model predictions
signature = infer_signature(train, predictions)
```
input_example – 一个或多个有效的模型输入实例。输入示例用作要馈送给模型的数据的提示。它将被转换为 Pandas DataFrame，然后使用 Pandas 的面向拆分（split-oriented）格式序列化为 json，或者转换为 numpy 数组，其中示例将通过转换为列表来序列化为 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.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.txt"）。如果提供，它将描述附加到根据用户当前软件环境自动生成的默认 pip 需求集之外的其他 pip 需求。需求和约束都会被自动解析并分别写入模型中的 requirements.txt 和 constraints.txt 文件。需求也会被写入模型 Conda 环境（conda.yaml）文件的 pip 部分。
警告

以下参数不能同时指定
- conda_env
- pip_requirements
- extra_pip_requirements
此示例演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip requirements。
save_model_kwargs – 默认为 None。传递给 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() 内部调用。

参数

model – keras.Model 的一个实例。要保存的 Keras 模型。
path – 要保存 MLflow 模型的本地路径。
save_exported_model – 如果为 True，则以导出的模型格式保存 Keras 模型，否则以 .keras 格式保存。有关更多信息，请参阅 https://keras.org.cn/guides/serialization_and_saving/。
conda_env –
Conda 环境的字典表示形式或 conda 环境 yaml 文件的路径。如果提供，它将描述此模型应运行的环境。至少，它应指定 get_default_conda_env() 中包含的依赖项。如果为 None，则将一个通过 mlflow.models.infer_pip_requirements() 推断出的 pip 需求的 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。MLflow 模型配置，Keras 模型元数据将被添加到其中。如果为 None，将创建一个空白实例。
signature –
一个 ModelSignature 类的实例，描述了模型的输入和输出。如果未指定但提供了 input_example，将根据提供的输入示例和模型自动推断签名。要在使用输入示例时禁用自动签名推断，请将 signature 设置为 False。要手动推断模型签名，请在具有有效模型输入的模型（例如，省略目标列的训练数据集）和有效模型输出（例如，在训练数据集上进行的模型预测）的数据集上调用 infer_signature()，例如：
```
from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # compute model predictions
signature = infer_signature(train, predictions)
```
input_example – 一个或多个有效的模型输入实例。输入示例用作要馈送给模型的数据的提示。它将被转换为 Pandas DataFrame，然后使用 Pandas 的面向拆分（split-oriented）格式序列化为 json，或者转换为 numpy 数组，其中示例将通过转换为列表来序列化为 json。字节将进行 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。
pip_requirements – pip 需求字符串的可迭代对象（例如 ["keras", "-r requirements.txt", "-c constraints.txt"]）或本地文件系统上的 pip 需求文件的字符串路径（例如 "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.txt"）。如果提供，它将描述附加到根据用户当前软件环境自动生成的默认 pip 需求集之外的其他 pip 需求。需求和约束都会被自动解析并分别写入模型中的 requirements.txt 和 constraints.txt 文件。需求也会被写入模型 Conda 环境（conda.yaml）文件的 pip 部分。
警告

以下参数不能同时指定
- conda_env
- pip_requirements
- extra_pip_requirements
此示例演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip requirements。
save_model_kwargs – 传递给 keras.Model.save 方法的 kwargs 字典。
metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

示例

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")