mlflow.tensorflow

The mlflow.tensorflow 模块提供用于记录和加载 TensorFlow 模型的 API。 该模块以以下 flavors 导出 TensorFlow 模型：

TensorFlow (native) format

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

mlflow.pyfunc

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

mlflow.tensorflow.autolog(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')[source]

注意

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。

Parameters

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

  • log_datasets – 如果 True，则将数据集信息记录到 MLflow Tracking。 如果 False，则不记录数据集信息。

  • disable – 如果 True，则禁用 TensorFlow 的 autologging 集成。如果 False，则启用 TensorFlow 的 autologging 集成。

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

  • disable_for_unsupported_versions – 如果 True，则为那些未经针对此版本的 MLflow 客户端测试或不兼容的 tensorflow 版本禁用自动记录（autologging）。

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

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

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

  • log_model_signatures – 如果 True，ModelSignatures 用于描述模型输入和输出，会在训练期间与 tf/keras 模型工件一起被收集并记录。如果 False，则不会记录签名。请注意，当向 predict() 传递 Pandas DataFrames 时，为带有签名的 TensorFlow 模型记录签名会改变它们的 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 – 如果为 True，则在每个 epoch 结束时记录训练指标。

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

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

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

  • checkpoint_mode – one of {"min", "max"}. 在自动模型检查点保存中， 如果 save_best_only=True，是否覆盖当前保存文件的决策基于对被监测指标的最大化或最小化。

  • checkpoint_save_best_only – If True, 自动模型检查点保存只有在根据所监控的量认为模型是“最佳”模型时才会进行，并且会覆盖之前的检查点模型。

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

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

mlflow.tensorflow.get_default_conda_env()[source]
Returns

默认的 Conda 环境，用于由调用 save_model() 和 log_model() 生成的 MLflow 模型。

mlflow.tensorflow.get_default_pip_requirements(include_cloudpickle=False)[source]
Returns

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

mlflow.tensorflow.get_global_custom_objects()[source]
Returns

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

mlflow.tensorflow.load_checkpoint(model=None, run_id=None, epoch=None, global_step=None)[source]

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

如果你想加载最新的检查点，请将 epoch 和 global_step 都设置为 None。 如果在 autologging 中将 “checkpoint_save_freq” 设置为 “epoch”，你可以将 epoch 参数设置为要加载的检查点的 epoch，以加载特定 epoch 的检查点。 如果在 autologging 中将 “checkpoint_save_freq” 设置为一个整数，你可以将 global_step 参数设置为要加载的检查点的 global step，以加载特定 global step 的检查点。 epoch 参数和 global_step 不能同时设置。

Parameters

  • model – 一个 Keras 模型，只有在保存的检查点为 “weight-only” 时才需要此参数。

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

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

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

Returns

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

Example 
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)[source]

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

Parameters

  • 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 方案的更多信息，请参见 Referencing Artifacts.

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

  • saved_model_kwargs – 要传递给 tensorflow.saved_model.load 方法的 kwargs。仅在加载 tensorflow2 core model 时可用。

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

Returns

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

mlflow.tensorflow.log_model(model, artifact_path: str | None = None, custom_objects=None, conda_env=None, code_paths=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, saved_model_kwargs=None, keras_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]

记录一个 TF2 核心模型（继承自 tf.Module）或一个以 MLflow Model 格式保存的 Keras 模型。

注意

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

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

Example of creating signature for saving TensorFlow and tf.Keras models 
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)
Parameters

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

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

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

  • conda_env –

    要么是 Conda 环境的字典表示，要么是指向 conda environment 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": [
                "tensorflow==x.y.z"
            ],
        },
    ],
}

  • code_paths –

    本地文件系统中指向 Python 文件依赖（或包含文件依赖的目录）路径的列表。这些文件在模型加载时会被预先添加到系统路径中。如果为某个模型声明了依赖文件且多个文件之间存在导入依赖关系，那么这些文件应从一个共同的根路径声明相对导入，以避免在加载模型时发生导入错误。

    有关 code_paths 功能、推荐的使用模式和限制的详细说明，请参阅 code_paths usage guide。

  • 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 – 如果提供，将在 registered_model_name 下创建一个模型版本，并在不存在同名的注册模型时创建该注册模型。

  • 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 文件的字符串路径（例如 "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 依赖。

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

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

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

  • name – 模型名称。

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

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

  • model_type – 模型的类型。

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

  • model_id – 模型的 ID。

Returns

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

mlflow.tensorflow.save_model(model, path, conda_env=None, code_paths=None, mlflow_model=None, custom_objects=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, saved_model_kwargs=None, keras_model_kwargs=None, metadata=None)[source]

将 TF2 core model（继承自 tf.Module）或 Keras model 以 MLflow Model 格式保存到本地文件系统中的某个路径。

注意

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

Example of creating signature for saving TensorFlow and tf.Keras models 
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)
Parameters

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

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

  • conda_env –

    可以是 Conda 环境的字典表示，或者是 conda 环境 yaml 文件的路径。如果提供，描述了应在其中运行此模型的环境。至少，它应指定包含在 get_default_conda_env() 中的依赖项。如果 None，将向模型添加一个 conda 环境，该环境包含由 mlflow.models.infer_pip_requirements() 推断的 pip 依赖项。如果依赖项推断失败，则回退使用 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 usage guide。

  • mlflow_model – MLflow 模型配置，将向其添加 tensorflow flavor。

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

  • 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 依赖字符串的可迭代对象 (e.g. ["tensorflow", "-r requirements.txt", "-c constraints.txt"]) 或者是指向本地文件系统上 pip 依赖文件的字符串路径 (e.g. "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 依赖。

  • saved_model_kwargs – 一个传递给 tensorflow.saved_model.save 方法的 kwargs 字典 如果要保存的模型是一个 Tensorflow 模块。

  • keras_model_kwargs – 一个字典，包含要传递给 model.save 方法的 kwargs，当要保存的模型是 keras 模型时使用。

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

class mlflow.tensorflow.MlflowCallback(log_every_epoch=True, log_every_n_steps=None)[source]

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

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

Parameters

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

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

Example 
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)[source]

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

on_epoch_end(epoch, logs=None)[source]

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

on_test_end(logs=None)[source]

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

on_train_begin(logs=None)[source]

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