跳到主要内容

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://<您的工作区主机>.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 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 回调​

如果您想添加额外的日志记录逻辑，您可以自定义 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 之间集成的相关指南。