跳至主内容

开始使用 MLflow + Tensorflow

下载此笔记本

在本指南中，我们将展示如何使用 Tensorflow 训练您的模型，并使用 MLflow 记录训练过程。

我们将使用 Databricks Free Trial，它对 MLflow 提供内置支持。Databricks Free Trial 提供了在 Databricks 平台上免费使用的机会，如果你还没有，请通过 link 注册一个账号。

您可以在本指南中通过云端笔记本（例如 Databricks notebook 或 Google Colab）运行代码，或在本地计算机上运行。

安装依赖​

让我们安装 mlflow 包。

%pip install -q mlflow

接下来让我们导入这些包。

import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow import keras

加载数据集​

我们将使用 mnist dataset 对手写数字进行一个简单的图像分类。

让我们使用 tensorflow_datasets（tfds）加载数据集，它以 tf.data.Dataset 的格式返回数据集。

# Load the mnist dataset.
train_ds, test_ds = tfds.load(
  "mnist",
  split=["train", "test"],
  shuffle_files=True,
)
Downloading and preparing dataset 11.06 MiB (download: 11.06 MiB, generated: 21.00 MiB, total: 32.06 MiB) to /root/tensorflow_datasets/mnist/3.0.1...
Dl Completed...:   0%|          | 0/5 [00:00<?, ? file/s]
Dataset mnist downloaded and prepared to /root/tensorflow_datasets/mnist/3.0.1. Subsequent calls will reuse this data.

让我们按以下步骤对数据进行预处理：

  • 将每个像素的值缩放到 [0, 1).
  • 将数据集分成批次。
  • 使用 prefetch 来加速训练。
def preprocess_fn(data):
  image = tf.cast(data["image"], tf.float32) / 255
  label = data["label"]
  return (image, label)


train_ds = train_ds.map(preprocess_fn).batch(128).prefetch(tf.data.AUTOTUNE)
test_ds = test_ds.map(preprocess_fn).batch(128).prefetch(tf.data.AUTOTUNE)

定义模型​

让我们将一个卷积神经网络定义为我们的分类器。我们可以使用 keras.Sequential 来堆叠这些层。

input_shape = (28, 28, 1)
num_classes = 10

model = keras.Sequential(
  [
      keras.Input(shape=input_shape),
      keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
      keras.layers.MaxPooling2D(pool_size=(2, 2)),
      keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
      keras.layers.MaxPooling2D(pool_size=(2, 2)),
      keras.layers.Flatten(),
      keras.layers.Dropout(0.5),
      keras.layers.Dense(num_classes, activation="softmax"),
  ]
)

设置训练相关的配置、优化器、损失函数、指标。

model.compile(
  loss=keras.losses.SparseCategoricalCrossentropy(),
  optimizer=keras.optimizers.Adam(0.001),
  metrics=[keras.metrics.SparseCategoricalAccuracy()],
)

设置跟踪/可视化工具​

在本教程中，我们将使用 Databricks 免费试用 作为 MLflow 跟踪服务器。有关其他选项，例如使用本地 MLflow 服务器，请阅读 Tracking Server Overview。

如果您还没有，请按照this guide设置您的 Databricks 免费试用帐户和访问令牌。注册应该不会超过 5 分钟。Databricks 免费试用是一种让用户免费试用 Databricks 功能的方式。对于本指南，我们需要 ML 实验仪表板来跟踪我们的训练进度。

在成功注册 Databricks 免费试用帐户后，让我们将 MLflow 连接到 Databricks 工作区。您需要输入以下信息：

  • Databricks 主机: https://.cloud.databricks.com/
  • Token: 你的个人访问令牌
import mlflow

mlflow.login()

现在这个 colab 已连接到托管的跟踪服务器。让我们配置 MLflow 元数据。需要设置两件事：

  • mlflow.set_tracking_uri: 始终使用 "databricks".
  • mlflow.set_experiment：选择你喜欢的名称，以 / 开头。

使用 MLflow 进行日志记录​

有两种方法可以将你的 Tensorflow 管道的日志记录到 MLflow：

  • MLflow 自动记录。
  • 使用回调。

自动记录配置简单，但提供的控制较少。使用回调更灵活。让我们看看每种方式是如何实现的。

MLflow 自动记录​

你只需在启动训练之前调用 mlflow.tensorflow.autolog()，然后后端会自动将指标记录到你之前配置的服务器上。在我们的例子中，为 Databricks Workspace。

# Choose any name that you like.
mlflow.set_experiment("/Users/<your email>/mlflow-tf-keras-mnist")

mlflow.tensorflow.autolog()

model.fit(x=train_ds, epochs=3)
2023/11/15 01:53:35 INFO mlflow.utils.autologging_utils: Created MLflow autologging run with ID '7c1db53e417b43f0a1d9e095c9943acb', which will track hyperparameters, performance metrics, model artifacts, and lineage information for the current tensorflow workflow
Epoch 1/3
469/469 [==============================] - 13s 7ms/step - loss: 0.3610 - sparse_categorical_accuracy: 0.8890
Epoch 2/3
469/469 [==============================] - 3s 6ms/step - loss: 0.1035 - sparse_categorical_accuracy: 0.9681
Epoch 3/3
469/469 [==============================] - 4s 8ms/step - loss: 0.0798 - sparse_categorical_accuracy: 0.9760
2023/11/15 01:54:05 WARNING mlflow.tensorflow: Failed to infer model signature: could not sample data to infer model signature: tuple index out of range
2023/11/15 01:54:05 WARNING mlflow.models.model: Model logged without a signature. Signatures will be required for upcoming model registry features as they validate model inputs and denote the expected schema of model outputs. Please visit https://www.mlflow.org/docs/2.8.1/models.html#set-signature-on-logged-model for instructions on setting a model signature on your logged model.
2023/11/15 01:54:05 WARNING mlflow.tensorflow: You are saving a TensorFlow Core model or Keras 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.
2023/11/15 01:54:13 WARNING mlflow.utils.autologging_utils: MLflow autologging encountered a warning: "/usr/local/lib/python3.10/dist-packages/_distutils_hack/__init__.py:33: UserWarning: Setuptools is replacing distutils."
Uploading artifacts:   0%|          | 0/11 [00:00<?, ?it/s]
2023/11/15 01:54:13 INFO mlflow.store.artifact.cloud_artifact_repo: The progress bar can be disabled by setting the environment variable MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR to false
Uploading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]
<keras.src.callbacks.History at 0x7d48e6556b60>

在训练进行期间，您可以在仪表板中找到该训练。登录到您的 Databricks Workspace，然后点击 Experiments tab。见下图： landing page

点击 Experiments 按钮后，你将进入实验页面，在那里可以找到你的运行。点击最新的实验和运行，你可以在那里找到你的指标，类似于： experiment page

您可以点击指标以查看图表。

让我们评估训练结果。

score = model.evaluate(test_ds)

print(f"Test loss: {score[0]:.4f}")
print(f"Test accuracy: {score[1]: .2f}")
79/79 [==============================] - 1s 12ms/step - loss: 0.0484 - sparse_categorical_accuracy: 0.9838
Test loss: 0.05
Test accuracy:  0.98

使用 MLflow 回调记录​

自动日志记录功能强大且方便，但如果您在寻找更原生的方式（例如 Tensorflow 的流水线），可以在 model.fit() 中使用 mlflow.tensorflow.MllflowCallback，它将记录：

  • 您的模型配置、层、超参数等。
  • 训练统计信息，包括使用 model.compile() 配置的损失和指标。
from mlflow.tensorflow import MlflowCallback

# Turn off autologging.
mlflow.tensorflow.autolog(disable=True)

with mlflow.start_run() as run:
  model.fit(
      x=train_ds,
      epochs=2,
      callbacks=[MlflowCallback(run)],
  )
Epoch 1/2
469/469 [==============================] - 5s 10ms/step - loss: 0.0473 - sparse_categorical_accuracy: 0.9851
Epoch 2/2
469/469 [==============================] - 4s 8ms/step - loss: 0.0432 - sparse_categorical_accuracy: 0.9866

进入 Databricks Workspace 的实验视图，您将看到与之前类似的仪表板。

自定义 MLflow Callback​

如果你想添加额外的日志记录逻辑，你可以自定义 MLflow 回调。你可以继承自 keras.callbacks.Callback 并从头编写所有内容，或者继承自 mlflow.tensorflow.MllflowCallback 来添加你自定义的日志记录逻辑。

让我们看一个示例，我们想将损失替换为其对数值并记录到 MLflow。

import math


# Create our own callback by subclassing `MlflowCallback`.
class MlflowCustomCallback(MlflowCallback):
  def on_epoch_end(self, epoch, logs=None):
      if not self.log_every_epoch:
          return
      loss = logs["loss"]
      logs["log_loss"] = math.log(loss)
      del logs["loss"]
      mlflow.log_metrics(logs, epoch)

使用新的回调训练模型。

with mlflow.start_run() as run:
  run_id = run.info.run_id
  model.fit(
      x=train_ds,
      epochs=2,
      callbacks=[MlflowCustomCallback(run)],
  )
Epoch 1/2
469/469 [==============================] - 5s 10ms/step - loss: 0.0537 - sparse_categorical_accuracy: 0.9834 - log_loss: -2.9237
Epoch 2/2
469/469 [==============================] - 4s 9ms/step - loss: 0.0497 - sparse_categorical_accuracy: 0.9846 - log_loss: -3.0022

进入你的 Databricks Workspace 页面，你应该会发现 log_loss 正在替代 loss 指标，类似于下面截图所示。

log loss screenshot

总结​

现在您已经了解了 MLflow 与 Tensorflow 之间的基本集成。本快速入门没有涵盖一些内容，例如将 TF 模型保存到 MLflow 以及将其加载回来的方法。有关详细指南，请参阅我们关于 MLflow 与 Tensorflow 集成的主指南。