mlflow.tensorflow

The mlflow.tensorflow 模块提供了一个用于记录和加载 TensorFlow 模型的 API。此模块导出具有以下**口味 (flavors)** 的 TensorFlow 模型：

TensorFlow（原生）格式: 这是可以加载回 TensorFlow 的主要格式。
mlflow.pyfunc: Produced for use by generic pyfunc-based deployment tools and batch inference.

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]

注意

自动日志记录已知与以下包版本兼容：2.16.0 <= tensorflow <= 2.20.0。在超出此范围的包版本中使用自动日志记录可能不成功。

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

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

tf.keras

指标 (Metrics) 和 参数 (Parameters)

训练和验证损失。

用户指定的指标。

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

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

工件

训练开始时的模型摘要。

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

训练结束时的 TensorBoard 日志。

tf.keras.callbacks.EarlyStopping

指标 (Metrics) 和 参数 (Parameters)

EarlyStopping 回调中的指标：stopped_epoch、restored_epoch、restore_best_weight 等

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

有关 TensorFlow 工作流的更多信息，请参阅自动日志记录跟踪文档。

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

参数

log_models – If True, trained models are logged as MLflow model artifacts. If False, trained models are not logged.
log_datasets – 如果为 True，则数据集信息将被记录到 MLflow Tracking。如果为 False，则不记录数据集信息。
disable – 如果为 True，则禁用 TensorFlow 自动日志记录集成。如果为 False，则启用 TensorFlow 集成自动日志记录集成。
exclusive – 如果为 True，则自动记录的内容不会记录到用户创建的流畅运行中。如果为 False，则自动记录的内容将记录到活动的流畅运行中，该运行可能是用户创建的。
disable_for_unsupported_versions – 如果为 True，则为尚未用此版本的 MLflow 客户端测试过或不兼容的 tensorflow 版本禁用自动日志记录。
silent – 如果为 True，则在 TensorFlow 自动日志记录期间抑制来自 MLflow 的所有事件日志和警告。如果为 False，则在 TensorFlow 自动日志记录期间显示所有事件和警告。
registered_model_name – If given, each time a model is trained, it is registered as a new model version of the registered model with this name. The registered model is created if it does not already exist.
log_input_examples – 如果为 True，则在训练期间收集并与 tf/keras 模型工件一起记录来自训练数据集的输入示例。如果为 False，则不记录输入示例。
log_model_signatures – 如果为 True，则在训练期间收集并与 tf/keras 模型工件一起记录描述模型输入和输出的 ModelSignatures**。如果为 False，则不记录签名。请注意，记录带有签名的 TensorFlow 模型会改变其 pyfunc 在向 predict() 传递 Pandas DataFrame 时的推理行为。当存在签名时，返回一个 np.ndarray（对于单输出模型）或 str -> np.ndarray 的映射（对于多输出模型）；当不存在签名时，返回一个 Pandas DataFrame。

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

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

extra_tags – 要为自动日志记录创建的每个托管运行设置的额外标签的字典。

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

log_every_n_steps – 如果设置，则每 n 个训练步骤记录一次训练指标。log_every_n_steps 必须为 None 才能使 log_every_epoch=True。

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

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

checkpoint_mode – {“min”, “max”} 中的一个。在自动模型检查点中，如果 save_best_only=True，则基于被监控数量的**最大化**或**最小化**来决定是否覆盖当前保存的文件。

checkpoint_save_best_only – 如果为 True，则自动模型检查点仅在模型根据被监控的数量被视为“最佳”模型时保存，并且会覆盖之前的检查点模型。

checkpoint_save_weights_only – 在自动模型检查点中，如果为 True，则只保存模型的权重。否则，检查点中还会添加优化器状态、lr-调度器状态等。

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

mlflow.tensorflow.get_default_conda_env()[source]

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

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

返回: 此格式由 save_model() 和 log_model() 生成的 MLflow 模型的默认 pip 要求列表。这些调用生成的 pip 环境至少包含这些要求。

mlflow.tensorflow.get_global_custom_objects()[source]

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

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

如果您在自动日志记录中启用了“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”，则要加载的检查点的 epoch。
global_step – 如果将“checkpoint_save_freq”设置为整数，则要加载的检查点的全局步骤。

返回

从指定检查点恢复的 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 格式的 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>
For more information about supported URI schemes, see Referencing Artifacts.
dst_path – The local filesystem path to which to download the model artifact. This directory must already exist. If unspecified, a local output path will be created.
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: 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]

以 MLflow Model 格式记录 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 – Deprecated. Use name instead.
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 要求。如果要求推断失败，它将回退到使用 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 –
A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the 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 = ...  # compute model predictions
signature = infer_signature(train, predictions)
```
input_example – 一个或多个有效的模型输入实例。输入示例用作要馈送给模型的数据的提示。它将被转换为 Pandas DataFrame，然后使用 Pandas 的面向拆分（split-oriented）格式序列化为 json，或者转换为 numpy 数组，其中示例将通过转换为列表来序列化为 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.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。
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。

返回: 一个 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 核心模型（继承自 tf.Module）或 Keras 模型以 MLflow 模型格式保存到本地文件系统的路径。

注意

If you save a Keras or TensorFlow model without a signature, inference with mlflow.pyfunc.spark_udf() will not work unless the model’s pyfunc representation accepts pandas DataFrames as inference inputs. You can infer a model’s signature by calling the mlflow.models.infer_signature() API on features from the model’s test dataset. You can also manually create a model signature, for example

示例：为保存 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 要求。如果要求推断失败，它将回退到使用 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 –
A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.
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()，例如
```
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 要求字符串的可迭代对象（例如 ["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 requirements。
saved_model_kwargs – 如果要保存的模型是 Tensorflow 模块，则传递给 tensorflow.saved_model.save 方法的 kwargs 字典。
keras_model_kwargs – 如果要保存的模型是 keras 模型，则传递给 model.save 方法的 kwargs 字典。
metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

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

用于将 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)[source]: 以用户指定的频率在每个批次结束时记录指标。

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

on_test_end(logs=None)[source]: 在验证结束时记录验证指标。

on_train_begin(logs=None)[source]: 在训练开始时记录模型架构和优化器配置。