由 Google 软件工程师 Jonah Kohn 和 Pavithra Vijay 发表

TensorFlow Cloud 是一个 Python 包，它提供 API 用于从在本地环境中调试和训练 TensorFlow 代码无缝过渡到在 Google Cloud 上进行分布式训练。它将云上训练模型的过程简化为一个简单的一次性函数调用，需要最少的设置，并且对您的模型几乎没有修改。TensorFlow Cloud 自动处理云相关任务，例如创建 VM 实例和为您的模型创建分布式策略。本文演示了 TensorFlow Cloud 的常见用例以及一些最佳实践。

我们将逐步介绍如何对 stanford_dogs 数据集提供的狗品种图像进行分类。为了简化操作，我们将使用在 ImageNet 权重上训练的 ResNet50 进行迁移学习。请在此处查看 TensorFlow Cloud 代码库中的此帖子的代码：此处。

设置

使用 pip install tensorflow_cloud 安装 TensorFlow Cloud。让我们从添加分类任务所需的导入开始编写 Python 脚本。

import datetime
import os

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_cloud as tfc
import tensorflow_datasets as tfds

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Model

Google Cloud 配置

TensorFlow Cloud 使用 Google Cloud 上的 AI Platform 服务在幕后运行您的训练作业。如果您不熟悉 GCP，请按照本节中的设置步骤创建和配置您的第一个 Google Cloud 项目。如果您不熟悉使用 Cloud，第一次设置和配置需要一些学习和工作。好消息是，设置完成后，您无需对 TensorFlow 代码进行任何更改即可在云上运行它！

1. 创建 GCP 项目
   
2. 启用 AI Platform 服务
   
3. 创建服务帐户
   
4. 下载授权密钥
   
5. 创建 Google Cloud Storage 存储桶
   

GCP 项目

Google Cloud 项目包含一组云资源，例如一组用户、一组 API、计费、身份验证和监控。要创建项目，请按照 本指南进行操作。在您的终端中运行本节中的命令。

export PROJECT_ID=<your-project-id>
gcloud config set project $PROJECT_ID

AI Platform 服务

请确保通过在 此下拉菜单中输入您的项目 ID 来为您的 GCP 项目启用 AI Platform 服务。

服务帐户和密钥

为您的新 GCP 项目创建 服务帐户。服务帐户是应用程序或 虚拟机实例使用的帐户，它由云应用程序用于进行授权的 API 调用。

export SA_NAME=<your-sa-name&rt;
gcloud iam service-accounts create $SA_NAME
gcloud projects add-iam-policy-binding $PROJECT_ID \
  --member serviceAccount:$SA_NAME@$PROJECT_ID.iam.gserviceaccount.com \
  --role 'roles/editor'

接下来，我们将需要服务帐户的身份验证密钥。此身份验证密钥是确保只有那些被授权处理您的项目的人员才会使用您的 GCP 资源的方法。按照以下步骤创建身份验证密钥

gcloud iam service-accounts keys create ~/key.json --iam-account $SA_NAME@$PROJECT_ID.iam.gserviceaccount.com

创建 GOOGLE_APPLICATION_CREDENTIALS 环境变量。

export GOOGLE_APPLICATION_CREDENTIALS=~/key.json

Cloud Storage 存储桶

如果您已经拥有指定的存储桶，请在下面输入您的存储桶名称。否则，请按照 本指南创建一个 Google Cloud Storage 存储桶。TensorFlow Cloud 使用 Google Cloud Build 来构建和发布 Docker 镜像，以及用于存储辅助数据（例如模型检查点和训练日志）。

GCP_BUCKET = "your-bucket-name"

Keras 模型创建

TensorFlow Cloud 的模型创建工作流程与在本地构建和训练 TF Keras 模型相同。

资源

我们将从加载 stanford_dogs 数据集开始，用于对狗品种进行分类。它作为 tensorflow-datasets 包的一部分提供。如果您有大型数据集，建议您将其托管在 GCS 上以获得更好的性能。

(ds_train, ds_test), metadata = tfds.load(
    "stanford_dogs",
    split=["train", "test"],
    shuffle_files=True,
    with_info=True,
    as_supervised=True,
)
 
NUM_CLASSES = metadata.features["label"].num_classes

让我们可视化数据集

print("Number of training samples: %d" % tf.data.experimental.cardinality(ds_train))
print("Number of test samples: %d" % tf.data.experimental.cardinality(ds_test))
print("Number of classes: %d" % NUM_CLASSES)

训练样本数量：12000 测试样本数量：8580 类别数量：120

plt.figure(figsize=(10, 10))
for i, (image, label) in enumerate(ds_train.take(9)):
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(image)
    plt.title(int(label))
    plt.axis("off")

预处理

我们将调整数据的大小并将其批处理。

IMG_SIZE = 224
BATCH_SIZE = 64
BUFFER_SIZE = 2
 
size = (IMG_SIZE, IMG_SIZE)
ds_train = ds_train.map(lambda image, label: (tf.image.resize(image, size), label))
ds_test = ds_test.map(lambda image, label: (tf.image.resize(image, size), label))
 
def input_preprocess(image, label):
    image = tf.keras.applications.resnet50.preprocess_input(image)
    return image, label

配置输入管道以提高性能

现在我们将配置输入管道以提高性能。请注意，我们正在使用并行调用和预取，这样 I/O 在您的模型训练期间不会成为阻塞。您可以在本 指南中了解有关配置输入管道以提高性能的更多信息。

ds_train = ds_train.map(
    input_preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE
)
 
ds_train = ds_train.batch(batch_size=BATCH_SIZE, drop_remainder=True)
ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)
 
ds_test = ds_test.map(input_preprocess)
ds_test = ds_test.batch(batch_size=BATCH_SIZE, drop_remainder=True)

构建模型

我们将加载在 ImageNet 上训练的 ResNet50，同时使用 include_top=False 以便为我们的任务重新整形模型。

inputs = tf.keras.layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
base_model = tf.keras.applications.ResNet50(
    weights="imagenet", include_top=False, input_tensor=inputs
)
x = tf.keras.layers.GlobalAveragePooling2D()(base_model.output)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(NUM_CLASSES)(x)
 
model = tf.keras.Model(inputs, outputs)

我们将冻结基本模型中的所有层及其当前权重，允许我们添加的额外层进行训练。

base_model.trainable = False

只要存储目的地位于您的 Cloud Storage 存储桶中，就可以轻松地在 TensorFlow Cloud 上使用 Keras 回调。在本例中，我们将使用 ModelCheckpoint 回调在训练的不同阶段保存模型，Tensorboard 回调可视化模型及其进度，以及 Early Stopping 回调来自动确定训练的最佳轮数。

MODEL_PATH = "resnet-dogs"
checkpoint_path = os.path.join("gs://", GCP_BUCKET, MODEL_PATH, "save_at_{epoch}")
tensorboard_path = os.path.join(
    "gs://", GCP_BUCKET, "logs", datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
)
callbacks = [
    tf.keras.callbacks.ModelCheckpoint(checkpoint_path),
    tf.keras.callbacks.TensorBoard(log_dir=tensorboard_path, histogram_freq=1),
    tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=3),
]

编译模型

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)

在本地调试模型

我们将首先在本地环境中训练模型，以确保代码在将作业发送到 GCP 之前正常运行。我们将使用 tfc.remote() 来确定代码应该在本地执行还是在云上执行。选择比预期的完整训练作业更少的轮数将有助于验证模型是否正常运行，而不会过载您的本地机器。

if tfc.remote():
    epochs = 500
    train_data = ds_train
    test_data = ds_test
else:
    epochs = 1
    train_data = ds_train.take(5)
    test_data = ds_test.take(5)
    callbacks = None
 
model.fit(
    train_data, epochs=epochs, callbacks=callbacks, validation_data=test_data, verbose=2
)

if tfc.remote():
    SAVE_PATH = os.path.join("gs://", GCP_BUCKET, MODEL_PATH)
    model.save(SAVE_PATH)

在 Google Cloud 上训练模型

要在 GCP 上进行训练，请在示例代码中填充您的 GCP 项目设置，然后只需从代码中调用 tfc.run()。API 很简单，对所有参数都有智能默认值。同样，在使用 TensorFlow Cloud 时，我们无需担心云相关任务，例如创建 VM 实例和分布式策略。依次，API 将

• 使您的 Python 脚本/笔记本云和分布式就绪。
  
• 将其转换为包含所需依赖项的 Docker 镜像。
  
• 在 GCP 集群上运行训练作业。
  
• 流式传输相关日志并存储检查点。
  

run() API 提供了显着的使用灵活性，例如允许用户指定自定义集群配置、自定义 Docker 镜像。有关可用于调用 run() 的所有参数的完整列表，请参阅 TensorFlow Cloud 自述文件。

创建一个 requirements.txt 文件，其中包含您的模型依赖的 Python 包列表。默认情况下，TensorFlow Cloud 将 TensorFlow 及其依赖项包含在默认 Docker 镜像中，因此无需包含这些依赖项。请在您的 Python 文件所在的同一目录中创建 requirements.txt。本例的 requirements.txt 内容为

tensorflow-datasets
matplotlib

默认情况下，run API 会根据您提供的集群配置将您的模型代码包装在 TensorFlow 分布式策略中。在本例中，我们使用的是单节点多 GPU 配置。因此，您的模型代码将自动包装在 TensorFlow `MirroredStrategy` 实例中。

调用 run() 以开始在云上进行训练。作业提交后，您将获得一个指向云作业的链接。要监控训练日志，请单击该链接并选择“查看日志”以查看训练进度信息。

tfc.run(
    requirements_txt="requirements.txt",
    distribution_strategy="auto",
    chief_config=tfc.MachineConfig(
        cpu_cores=8,
        memory=30,
        accelerator_type=tfc.AcceleratorType.NVIDIA_TESLA_T4,
        accelerator_count=2,
    ),
    docker_image_bucket_name=GCP_BUCKET,
)

使用 TensorBoard 可视化模型

在这里，我们正在从 GCS 存储桶加载 TensorBoard 日志以评估模型性能和历史记录。

tensorboard dev upload --logdir "gs://your-bucket-name/logs" --name "ResNet Dogs"

评估模型

训练完成后，我们可以加载存储在 GCS 存储桶中的模型，并评估其性能。

if tfc.remote():
    model = tf.keras.models.load_model(SAVE_PATH)
model.evaluate(test_data)

后续步骤

本文介绍了 TensorFlow Cloud，它是一个 Python 包，它简化了使用多个 GPU/TPU 在云上训练模型的过程，只需一个函数即可完成，而无需对您的模型进行任何代码更改。您可以在此处找到本文的完整代码：此处。作为下一步，您可以在 TensorFlow Cloud 代码库中找到此代码示例以及许多其他示例。