mlflow.tensorflow

mlflow.tensorflow 模块提供了一个用于记录和加载 TensorFlow 模型的 API。该模块以以下格式导出 TensorFlow 模型：

TensorFlow (原生) 格式

这是可以重新加载到 TensorFlow 中的主要风格。

mlflow.pyfunc

为基于通用 pyfunc 的部署工具和批量推理而生成。

mlflow.tensorflow.autolog(every_n_iter=1, log_models=True, log_datasets=True, disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False, registered_model_name=None, log_input_examples=False, log_model_signatures=True, saved_model_kwargs=None, keras_model_kwargs=None, extra_tags=None, log_every_epoch=True, log_every_n_steps=None, checkpoint=True, checkpoint_monitor='val_loss', checkpoint_mode='min', checkpoint_save_best_only=True, checkpoint_save_weights_only=False, checkpoint_save_freq='epoch')[源代码]

备注

Autologging 已知与以下包版本兼容：2.7.4 <= tensorflow <= 2.17.0。当与该范围之外的包版本一起使用时，Autologging 可能不会成功。

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

对于每个 TensorFlow 模块，自动记录捕获以下信息：

tf.keras

  • 指标 和 参数

  • 训练和验证损失。

  • 用户指定的指标。

  • 优化器配置，例如，学习率、动量等。

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

  • 工件

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

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

  • 训练结束时的 TensorBoard 日志。

tf.keras.callbacks.EarlyStopping

  • 指标 和 参数

  • 来自 EarlyStopping 回调的指标：stopped_epoch、restored_epoch、restore_best_weight 等

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

有关 TensorFlow 工作流 的更多信息，请参阅 autologging 跟踪文档。

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

参数:

  • every_n_iter – 已弃用，请改用 log_every_epoch。每 every_n_iter 步，将记录指标。

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

  • log_datasets – 如果 True，数据集信息将被记录到 MLflow 跟踪中。如果 False，数据集信息将不会被记录。

  • disable – 如果 True，禁用 TensorFlow 自动日志集成。如果 False，启用 TensorFlow 自动日志集成。

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

  • disable_for_unsupported_versions – 如果 True，则对未经过此版本 MLflow 客户端测试或不兼容的 TensorFlow 版本禁用自动日志记录。

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

  • registered_model_name – 如果提供，每次训练模型时，它都会被注册为具有此名称的已注册模型的新的模型版本。如果该注册模型尚不存在，则会创建它。

  • log_input_examples – 如果 True，训练数据集中的输入示例将与训练期间的 tf/keras 模型工件一起收集和记录。如果 False，则不记录输入示例。

  • log_model_signatures – 如果 True，在训练期间会收集并记录描述模型输入和输出的 模型签名 以及 tf/keras 模型工件。如果 False，则不记录签名。请注意，使用签名记录 TensorFlow 模型会改变当 Pandas DataFrame 传递给 predict() 时的 pyfunc 推理行为。当存在签名时，返回 np.ndarray``（对于单输出模型）或从 ``str 到 np.ndarray 的映射（对于多输出模型）；当不存在签名时，返回 Pandas DataFrame。

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

  • keras_model_kwargs – 传递给 keras_model.save 方法的 kwargs 字典。

  • extra_tags – 一个字典，包含要为 autologging 创建的每个托管运行设置的额外标签。

  • log_every_epoch – 如果为真，训练指标将在每个 epoch 结束时记录。

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

  • checkpoint – 启用自动模型检查点保存。

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

  • checkpoint_save_best_only – 如果为 True，自动模型检查点仅在模型被认为是根据监控量和之前的检查点模型被覆盖的“最佳”模型时保存。

  • checkpoint_mode – one of {“min”, “max”}。在自动模型检查点保存中，如果 save_best_only=True，则根据监视量的最大化或最小化来决定是否覆盖当前保存文件。

  • checkpoint_save_weights_only – 在自动模型检查点保存中，如果为 True，则仅保存模型的权重。否则，优化器状态、学习率调度器状态等也会添加到检查点中。

  • checkpoint_save_freq – “epoch” 或整数。当使用 “epoch” 时，回调函数在每个 epoch 后保存模型。当使用整数时，回调函数在此数量的批次结束时保存模型。请注意，如果保存与 epoch 不对齐，监控的指标可能不太可靠（因为它可能只反映一个批次，因为指标在每个 epoch 都会重置）。默认为 “epoch”。

mlflow.tensorflow.get_default_conda_env()[源代码]
返回:

通过调用 save_model() 和 log_model() 生成的 MLflow 模型的默认 Conda 环境。

mlflow.tensorflow.get_default_pip_requirements(include_cloudpickle=False)[源代码]
返回

此flavor生成的MLflow Models的默认pip需求列表。对 save_model() 和 log_model() 的调用会生成一个pip环境，该环境至少包含这些需求。

mlflow.tensorflow.get_global_custom_objects()[源代码]
返回:

对自定义对象全局字典的实时引用。

mlflow.tensorflow.load_checkpoint(model=None, run_id=None, epoch=None, global_step=None)[源代码]

如果在 autologging 中启用 “checkpoint”，在 Keras 模型训练执行期间，检查点模型会被记录为 MLflow 工件。使用此 API，您可以加载检查点模型。

如果你想加载最新的检查点，将 epoch 和 global_step 都设置为 None。如果在自动记录中将 checkpoint_save_freq 设置为 epoch，你可以将 epoch 参数设置为要加载的检查点的 epoch 来加载特定 epoch 的检查点。如果在自动记录中将 checkpoint_save_freq 设置为一个整数，你可以将 global_step 参数设置为要加载的检查点的全局步数来加载特定全局步数的检查点。epoch 参数和 global_step 不能同时设置。

参数:

  • model – Keras 模型，只有当保存的检查点是“仅权重”时，此参数才是必需的。

  • run_id – 模型日志记录到的运行的ID。如果未提供，则使用当前活动的运行。

  • epoch – 要加载的检查点的时期，如果你将 “checkpoint_save_freq” 设置为 “epoch”。

  • global_step – 如果要加载的检查点的全局步骤，如果你将“checkpoint_save_freq”设置为一个整数。

返回:

从指定检查点恢复的 Keras 模型实例。

示例 
import mlflow

mlflow.tensorflow.autolog(checkpoint=True, checkpoint_save_best_only=False)

model = create_tf_keras_model()  # Create a Keras model
with mlflow.start_run() as run:
    model.fit(data, label, epoch=10)

run_id = run.info.run_id

# load latest checkpoint model
latest_checkpoint_model = mlflow.tensorflow.load_checkpoint(run_id=run_id)

# load history checkpoint model logged in second epoch
checkpoint_model = mlflow.tensorflow.load_checkpoint(run_id=run_id, epoch=2)
mlflow.tensorflow.load_model(model_uri, dst_path=None, saved_model_kwargs=None, keras_model_kwargs=None)[源代码]

从指定路径加载包含 TensorFlow 风格的 MLflow 模型。

参数:

  • model_uri – MLflow 模型的位置，采用 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> 有关支持的 URI 方案的更多信息，请参阅 引用工件。

  • dst_path – 下载模型工件的本地文件系统路径。此目录必须已经存在。如果未指定，将创建一个本地输出路径。

  • saved_model_kwargs – 传递给 tensorflow.saved_model.load 方法的 kwargs。仅在加载 tensorflow2 核心模型时可用。

  • keras_model_kwargs – 传递给 keras.models.load_model 方法的 kwargs。仅在你加载 Keras 模型时可用。

返回

一个可调用的图（tf.function），它接受输入并返回推理结果。

mlflow.tensorflow.log_model(model, artifact_path, custom_objects=None, conda_env=None, code_paths=None, signature: ModelSignature = None, input_example: DataFrame | 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, saved_model_kwargs=None, keras_model_kwargs=None, metadata=None)[源代码]

在 MLflow 模型格式中记录一个 TF2 核心模型（继承自 tf.Module）或一个 Keras 模型。

备注

如果你在没有签名的前提下记录一个 Keras 或 TensorFlow 模型，使用 mlflow.pyfunc.spark_udf() 进行推理将无法工作，除非模型的 pyfunc 表示接受 pandas DataFrame 作为推理输入。

你可以通过调用 mlflow.models.infer_signature() API 从模型的测试数据集中推断模型的签名。你也可以手动创建模型签名，例如：

创建用于保存 TensorFlow 和 tf.Keras 模型的签名的示例 
from mlflow.types.schema import Schema, TensorSpec
from mlflow.models import ModelSignature
import numpy as np

input_schema = Schema(
    [
        TensorSpec(np.dtype(np.uint64), (-1, 5), "field1"),
        TensorSpec(np.dtype(np.float32), (-1, 3, 2), "field2"),
    ]
)
# Create the signature for a model that requires 2 inputs:
#  - Input with name "field1", shape (-1, 5), type "np.uint64"
#  - Input with name "field2", shape (-1, 3, 2), type "np.float32"
signature = ModelSignature(inputs=input_schema)
参数:

  • model – 要保存的 TF2 核心模型（继承自 tf.Module）或 Keras 模型。

  • artifact_path – 记录模型工件的运行相对路径。

  • custom_objects – 一个 Keras custom_objects 字典，将名称（字符串）映射到与 Keras 模型相关联的自定义类或函数。MLflow 使用 CloudPickle 保存这些自定义层，并在使用 mlflow.tensorflow.load_model() 和 mlflow.pyfunc.load_model() 加载模型时自动恢复它们。

  • 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": [
                "tensorflow==x.y.z"
            ],
        },
    ],
}

  • code_paths – 本地文件系统路径列表，指向Python文件依赖项（或包含文件依赖项的目录）。这些文件在加载模型时会被*前置*到系统路径中。如果为给定模型声明了依赖关系，则应从公共根路径声明相对导入，以避免在加载模型时出现导入错误。有关``code_paths``功能的详细解释、推荐的使用模式和限制，请参阅`code_paths使用指南 <https://mlflow.org/docs/latest/model/dependencies.html?highlight=code_paths#saving-extra-code-with-an-mlflow-model>`_。

  • registered_model_name – 如果指定，在 registered_model_name 下创建一个模型版本，如果给定名称的注册模型不存在，则同时创建一个注册模型。

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

  • input_example – 一个或多个有效的模型输入实例。输入示例用作提示，指示应向模型提供哪些数据。它将被转换为Pandas DataFrame，然后使用Pandas的面向分割的格式序列化为json，或者是一个numpy数组，其中示例将通过将其转换为列表来序列化为json。字节被base64编码。当``signature``参数为``None``时，输入示例用于推断模型签名。

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

  • pip_requirements – 可以是 pip 需求字符串的可迭代对象（例如 ["tensorflow", "-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 需求。

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

  • keras_model_kwargs – 传递给 keras_model.save 方法的 kwargs 字典。

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

返回

一个包含记录模型元数据的 ModelInfo 实例。

mlflow.tensorflow.save_model(model, path, conda_env=None, code_paths=None, mlflow_model=None, custom_objects=None, signature: ModelSignature = None, input_example: DataFrame | ndarray | dict | list | csr_matrix | csc_matrix | str | bytes | tuple = None, pip_requirements=None, extra_pip_requirements=None, saved_model_kwargs=None, keras_model_kwargs=None, metadata=None)[源代码]

将 TF2 核心模型（继承自 tf.Module）或 Keras 模型以 MLflow 模型格式保存到本地文件系统上的路径。

备注

如果你保存了一个没有签名的 Keras 或 TensorFlow 模型，使用 mlflow.pyfunc.spark_udf() 进行推理将无法工作，除非模型的 pyfunc 表示接受 pandas DataFrame 作为推理输入。你可以通过在模型的测试数据集的特征上调用 mlflow.models.infer_signature() API 来推断模型的签名。你也可以手动创建一个模型签名，例如：

创建用于保存 TensorFlow 和 tf.Keras 模型的签名的示例 
from mlflow.types.schema import Schema, TensorSpec
from mlflow.models import ModelSignature
import numpy as np

input_schema = Schema(
    [
        TensorSpec(np.dtype(np.uint64), (-1, 5), "field1"),
        TensorSpec(np.dtype(np.float32), (-1, 3, 2), "field2"),
    ]
)
# Create the signature for a model that requires 2 inputs:
#  - Input with name "field1", shape (-1, 5), type "np.uint64"
#  - Input with name "field2", shape (-1, 3, 2), type "np.float32"
signature = ModelSignature(inputs=input_schema)
参数:

  • model – 要保存的 Keras 模型或 Tensorflow 模块。

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

  • 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": [
                "tensorflow==x.y.z"
            ],
        },
    ],
}

  • code_paths – 本地文件系统路径列表，指向Python文件依赖项（或包含文件依赖项的目录）。这些文件在加载模型时会被*前置*到系统路径中。如果为给定模型声明了依赖关系，则应从公共根路径声明相对导入，以避免在加载模型时出现导入错误。有关``code_paths``功能的详细解释、推荐的使用模式和限制，请参阅`code_paths使用指南 <https://mlflow.org/docs/latest/model/dependencies.html?highlight=code_paths#saving-extra-code-with-an-mlflow-model>`_。

  • mlflow_model – 要添加 tensorflow 风格的 MLflow 模型配置。

  • custom_objects – 一个 Keras custom_objects 字典，将名称（字符串）映射到与 Keras 模型相关联的自定义类或函数。MLflow 使用 CloudPickle 保存这些自定义层，并在使用 mlflow.tensorflow.load_model() 和 mlflow.pyfunc.load_model() 加载模型时自动恢复它们。

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

  • input_example – 一个或多个有效的模型输入实例。输入示例用作提示，指示应向模型提供哪些数据。它将被转换为Pandas DataFrame，然后使用Pandas的面向分割的格式序列化为json，或者是一个numpy数组，其中示例将通过将其转换为列表来序列化为json。字节被base64编码。当``signature``参数为``None``时，输入示例用于推断模型签名。

  • pip_requirements – 可以是 pip 需求字符串的可迭代对象（例如 ["tensorflow", "-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 需求。

  • saved_model_kwargs – 一个字典，包含传递给 tensorflow.saved_model.save 方法的关键字参数，如果需要保存的模型是一个 Tensorflow 模块。

  • keras_model_kwargs – 一个字典，包含要传递给 model.save 方法的关键字参数，如果需要保存的模型是 Keras 模型。

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

class mlflow.tensorflow.MlflowCallback(log_every_epoch=True, log_every_n_steps=None)[源代码]

用于将 Tensorflow 训练指标记录到 MLflow 的回调函数。

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

参数:

  • log_every_epoch – bool, 如果为 True，则每个 epoch 记录一次指标。如果为 False，则每 n 步记录一次指标。

  • log_every_n_steps – int, 每 n 步记录一次指标。如果为 None，则每个 epoch 记录一次指标。如果 log_every_epoch=True，则必须为 None。

示例 
from tensorflow import keras
import mlflow
import numpy as np

# Prepare data for a 2-class classification.
data = tf.random.uniform([8, 28, 28, 3])
label = tf.convert_to_tensor(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(run)],
    )
on_batch_end(batch, logs=None)[源代码]

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

on_epoch_end(epoch, logs=None)[源代码]

在每个epoch结束时记录指标。

on_test_end(logs=None)[源代码]

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

on_train_begin(logs=None)[源代码]

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