跳到主要内容

MLflow + Tensorflow 入门

下载此笔记本

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

我们将使用 Databricks 免费试用版，其中内置了对 MLflow 的支持。Databricks 免费试用版提供了一个免费使用 Databricks 平台的机会。如果您还没有注册，请通过 此链接注册一个帐户。

您可以从基于云的笔记本（如 Databricks 笔记本或 Google Colab）运行本指南中的代码，或者在本地计算机上运行它。

安装依赖项​

让我们安装 mlflow 包。

%pip install -q mlflow

然后，我们导入必要的包。

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

加载数据集​

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

让我们使用 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 服务器，请阅读 跟踪服务器概述。

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

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

  • Databricks 主机: https://<your workspace host>.cloud.databricks.com/
  • 令牌：您的个人访问令牌
import mlflow

mlflow.login()

现在，此 Colab 已连接到托管的跟踪服务器。让我们配置 MLflow 元数据。需要设置两项内容：

  • mlflow.set_tracking_uri：始终使用“databricks”。
  • mlflow.set_experiment：选择一个您喜欢的名称，以 / 开头。

使用 MLflow 进行日志记录​

您可以通过两种方式从 Tensorflow 管道将数据记录到 MLflow：

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

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

MLflow 自动日志记录​

您需要做的就是在开始训练之前调用 mlflow.tensorflow.autolog()，然后后端会自动将指标记录到您之前配置的服务器中。在本例中，是 Databricks 工作区。

# 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 工作区，然后单击 Experiments 选项卡。请参阅下面的截图： 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 工作区的实验视图，您将看到与之前类似的仪表板。

自定义 MLflow 回调​

如果您想添加额外的日志记录逻辑，可以自定义 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 工作区页面，您应该会发现 log_loss 替换了 loss 指标，类似于下面的屏幕截图。

log loss screenshot

总结​

现在您已经了解了 MLflow 和 Tensorflow 的基本集成。此快速入门指南并未涵盖所有内容，例如将 TF 模型保存到 MLflow 并将其加载回来。有关详细指南，请参阅我们关于 MLflow 和 Tensorflow 集成的主要指南。