跳到主要内容

MLflow + Tensorflow 入门

下载此 Notebook

本指南将展示如何使用 Tensorflow 训练模型并使用 MLflow 记录训练过程。

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

您可以在基于云的 Notebook 中运行本指南中的代码，例如 Databricks Notebook 或 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://<您的 workspace 主机>.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，然后单击“实验”选项卡。请参阅下面的屏幕截图： 主页

点击“实验”按钮后，您将进入实验页面，在那里可以找到您的运行。点击最近的实验和运行，您可以在那里找到您的指标，类似于： 实验页

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

让我们评估训练结果。

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 集成主指南。