TF2 Keras (13) : 在Tensorflow Clould 上训练Keras 模型

本文是对官方文档 的学习笔记。

---

TensorFlow Cloud是一个Python软件包，提供了API，可实现从本地调试到Google Cloud中分布式培训的无缝过渡。它简化了将云上的TensorFlow模型训练为一个简单的函数调用的过程，只需最少的设置即可完成，而无需更改模型。 TensorFlow Cloud可处理特定于云的任务，例如自动为模型创建VM实例和分发策略。本指南将演示如何通过TensorFlow Cloud与Google Cloud进行接口，以及TensorFlow Cloud中提供的广泛功能。我们将从最简单的用例开始。

准备工作

安装相关包

pip install -q tensorflow_cloud

引入包

import tensorflow as tf
import tensorflow_cloud as tfc

from tensorflow import keras
from tensorflow.keras import layers

样例

本例是一个CNN 模型

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

model = keras.Sequential(
    [
        keras.Input(shape=(28, 28)),
        # Use a Rescaling layer to make sure input values are in the [0, 1] range.
        layers.experimental.preprocessing.Rescaling(1.0 / 255),
        # The original images have shape (28, 28), so we reshape them to (28, 28, 1)
        layers.Reshape(target_shape=(28, 28, 1)),
        # Follow-up with a classic small convnet
        layers.Conv2D(32, 3, activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(32, 3, activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(32, 3, activation="relu"),
        layers.Flatten(),
        layers.Dense(128, activation="relu"),
        layers.Dense(10),
    ]
)

model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=keras.metrics.SparseCategoricalAccuracy(),
)

model.fit(x_train, y_train, epochs=20, batch_size=128, validation_split=0.1)

调用run（） 使得模型在 Google Cloud 上进行训练

tfc.run()

使用TensorFlow Cloud时，无需担心特定于云的任务，例如创建VM实例和分发策略。该API包括所有参数的智能默认值-一切都是可配置的，但是许多模型都可以依赖这些默认值。

调用run（）后，TensorFlow Cloud将：

• 使Python脚本或笔记本做好发布准备。
• 将其转换为具有所需依赖项的Docker映像。
• 在GCP GPU驱动的VM上训练。
• 传递相关日志和作业信息。
• 默认的VM配置为1主和0从，具有8个CPU内核和1个Tesla T4 GPU。

Google Cloud 配置

为了进行Cloud Training，需要进行一些首次设置。如果是Google Cloud的新用户，则需要采取一些初步步骤：

• 创建一个GCP项目
• 启用AI平台服务
• 创建一个服务帐户
• 下载授权密钥
• 创建一个云存储 bucket

详细信息可以查看 TensorFlow Cloud README
更多的例子可以参考 TensorFlow Blog.

常用工作流和云存储

在大多数情况下，我们都要在训练完毕后从云上取下模型。为此，在远程 训练 时将保存和加载重定向到Cloud Storage至关重要。我们可以将TensorFlow Cloud 定向到我们的Cloud Storage存储 以执行各种任务。 云存储可用于保存和加载大型训练数据集，存储回调日志或模型权重以及保存训练后的模型文件。首先，让我们配置fit（）以将模型保存到Cloud Storage，并设置TensorBoard监控以跟踪训练进度。

def create_model():
    model = keras.Sequential(
        [
            keras.Input(shape=(28, 28)),
            layers.experimental.preprocessing.Rescaling(1.0 / 255),
            layers.Reshape(target_shape=(28, 28, 1)),
            layers.Conv2D(32, 3, activation="relu"),
            layers.MaxPooling2D(2),
            layers.Conv2D(32, 3, activation="relu"),
            layers.MaxPooling2D(2),
            layers.Conv2D(32, 3, activation="relu"),
            layers.Flatten(),
            layers.Dense(128, activation="relu"),
            layers.Dense(10),
        ]
    )

    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=keras.metrics.SparseCategoricalAccuracy(),
    )
    return model

将 TensorBoard logs， checkpoints 存储到云存储中：

import datetime
import os

# Note: Please change the gcp_bucket to your bucket name.
gcp_bucket = "keras-examples"

checkpoint_path = os.path.join("gs://", gcp_bucket, "mnist_example", "save_at_{epoch}")

tensorboard_path = os.path.join(  # Timestamp included to enable timeseries graphs
    "gs://", gcp_bucket, "logs", datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
)

callbacks = [
    # TensorBoard will store logs for each epoch and graph performance for us.
    keras.callbacks.TensorBoard(log_dir=tensorboard_path, histogram_freq=1),
    # ModelCheckpoint will save models after each epoch for retrieval later.
    keras.callbacks.ModelCheckpoint(checkpoint_path),
    # EarlyStopping will terminate training when val_loss ceases to improve.
    keras.callbacks.EarlyStopping(monitor="val_loss", patience=3),
]

model = create_model()

在这里，我们将直接从Keras加载数据。通常，最佳做法是将数据集存储在Cloud Storage存储桶中，但是TensorFlow Cloud也可以适应本地存储的数据集。本指南的“多文件”部分对此进行了介绍。

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

TensorFlow Cloud API提供了remote（）函数来确定代码是在本地执行还是在云上执行。这样就可以为本地和远程执行分别指定fit（）参数，并提供简便的调试方法而又不会使本地计算机过载。

if tfc.remote():
    epochs = 100
    callbacks = callbacks
    batch_size = 128
else:
    epochs = 5
    batch_size = 64
    callbacks = None

model.fit(x_train, y_train, epochs=epochs, callbacks=callbacks, batch_size=batch_size)

在训练结束后， 将模型存储在GCS

save_path = os.path.join("gs://", gcp_bucket, "mnist_example")

if tfc.remote():
    model.save(save_path)

我们还可以使用此存储桶来构建Docker映像，而不是使用本地Docker实例。为此，只需将您的云储存添加到docker_image_bucket_name参数。

tfc.run(docker_image_bucket_name=gcp_bucket)

load modle 并且在 TensorBoard 监控性能

model = keras.models.load_model(save_path)

!tensorboard dev upload --logdir "gs://keras-examples-jonah/logs/fit" --name "Guide MNIST"

大型项目

在许多情况下，包含Keras模型的项目可能包含多个Python脚本，或者涉及外部数据或特定的依赖项。 TensorFlow Cloud对于大规模部署具有完全的灵活性，并提供了许多智能功能来帮助您的项目。

Entry points: 对 Python scripts 和Jupyter notebooks 的支持

对run（）API的调用不会总是包含在与模型训练代码相同的Python脚本中。为此，我们提供了一个entry_point参数。 entry_point参数可用于指定模型训练代码所在的Python脚本或笔记本。从与模型相同的脚本中调用run（）时，请使用entry_point的默认值None。

pip 依赖

如果您的项目需要其他pip依赖项，则可以通过包含requirements.txt文件来指定其他所需的库。在此文件中，只需列出所有必需依赖项的列表，TensorFlow Cloud便会处理将这些依赖项集成到您的云构建中。

Python notebooks

TensorFlow Cloud也可以从Python notebooks运行。此外，如果需要，指定的entry_point可以是notebooks。与脚本相比，notebooks 上的TensorFlow Cloud要记住两个主要区别：

• 从notebooks 调用 run 时， 必须指定云存储用的bucket
• Google Cloud 验证是对所有 project

多文件项目

如果模型依赖于其他文件，则只需确保这些文件位于指定入口点的同一目录（或子目录）中。与指定的entry_point存储在同一目录中的每个文件，以及与entry_point相邻的子目录中存储的任何文件，都将包含在Docker映像中。对于可能需要通过pip无法获得的依赖项也是如此

有关具有其他pip依赖项的自定义入口点和多文件项目的示例，请在 TensorFlow Cloud Repository. 为简便起见，仅包含示例的run（）调用：

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    entry_point="train_model.py",
    requirements="requirements.txt"
)

机器配置与分布式训练

模型训练可能需要不同资源，具体取决于模型或数据集的大小。考虑具有多个GPU的配置时，选择合适的分配策略就变得至关重要。在这里，我们概述了一些可能的配置：

多 Worker 配置

通过 COMMON_MACHINE_CONFIGS ，配置为 1CPU +4 GPU

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    chief_config=tfc.COMMON_MACHINE_CONFIGS['CPU'],
    worker_count=2,
    worker_config=tfc.COMMON_MACHINE_CONFIGS['T4_4X']
)

默认情况下，TensorFlow Cloud使用提供的Chief_config，worker_config和worker_count参数的简单公式为机器配置选择最佳的分发策略。

• 如果指定的GPU数量大于零，则将选择tf.distribute.MirroredStrategy。
• 如果工作者数量大于零，则将基于加速器类型选择tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.experimental.TPUStrategy。
• 否则，将选择tf.distribute.OneDeviceStrategy。

TPU distribution

让我们在TPU上训练相同的模型，如下所示：

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    chief_config=tfc.COMMON_MACHINE_CONFIGS["CPU"],
    worker_count=1,
    worker_config=tfc.COMMON_MACHINE_CONFIGS["TPU"]
)

定制化分布式策略

如果想自定义分布式策略， 可以参考 distributed training guide ， 并且将 distribution_strategy 设置为 None 。 下面的例子是一个针对 MINST 的自定义分布式策略的：

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
  model = create_model()

if tfc.remote():
    epochs = 100
    batch_size = 128
else:
    epochs = 10
    batch_size = 64
    callbacks = None

model.fit(
    x_train, y_train, epochs=epochs, callbacks=callbacks, batch_size=batch_size
)

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    chief_config=tfc.COMMON_MACHINE_CONFIGS['CPU'],
    worker_count=2,
    worker_config=tfc.COMMON_MACHINE_CONFIGS['T4_4X'],
    distribution_strategy=None
)

自定义 Docker Image

默认情况下，TensorFlow Cloud使用Google提供的与您当前TensorFlow版本相对应的Docker基本映像。但是，如果需要，您还可以指定一个自定义Docker映像来满足您的构建要求。在此示例中，我们将从旧版本的TensorFlow中指定Docker映像：

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    base_docker_image="tensorflow/tensorflow:2.1.0-gpu"
)

附加指标

您可能会发现用特定标签标记Cloud作业或在Cloud培训期间流式传输模型日志很有用。好的作法是在所有Cloud作业上保持适当的标签，以保持记录。为此，run（）接受标签字典，最多可包含64个键值对，这些标签可从Cloud构建日志中看到。可以使用执行tfc.run所提供的链接来访问诸如时代性能和模型保存内部的日志，或使用stream_logs标志将其打印到本地终端。

job_labels = {"job": "mnist-example", "team": "keras-io", "user": "jonah"}

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    job_labels=job_labels,
    stream_logs=True
)

综合例子

对于使用本指南中描述的许多功能的深入Colab，请按照本示例进行操作，以训练一个最新的模型，以使用特征提取从照片中识别狗的品种。

this example