教程 | 在Cloud ML Engine的TPU上从头训练ResNet

选自GoogleCloud

作者：Lak Lakshmanan

机器之心编译

参与：Geek AI、王淑婷

以往的测试显示，张量处理单元（TPU）是能够极大加快深度学习模型训练速度的存在。本文作者将演示如何使用谷歌云提供的 TPU 在自己的数据集上训练一个最先进的图像分类模型。文中还包含了详细的教程目录和内容，心动的读者不妨跟着一起动手试试？

张量处理单元（TPU）是能够大大加快深度学习模型训练速度的硬件加速器。在斯坦福大学进行的独立测试中，在 TPU 上训练的 ResNet-50 模型能够在 ImageNet 数据集上以最快的速度（30 分钟）达到预期的准确率。

在本文中，我将带领读者使用谷歌云提供的 TPU 在自己的数据集上训练一个最先进的图像分类模型。并且：

• 无需自行编写 TensorFlow 代码（我已经完成了所有代码。）
• 不需要安装软件或基础环境（Cloud ML Engine 是无服务器的）
• 你可以在云端训练模型，然后在任何地方部署该模型（使用 Kubeflow）
• 作者写的代码：https://github.com/tensorflow/tpu/tree/master/models/official/resnet
• Cloud ML Engine：https://cloud.google.com/ml-engine/docs/tensorflow/technical-overview
• Kubeflow：https://github.com/kubeflow/kubeflow

完整的代码存放在 GitHub 的一个 notebook 中。读者可以使用这个 notebook 或这个 codelab 中的代码来跟进此教程。我已经在 Cloud Datalab 中测试了 notebook，并且在 Cloud Shell 中测试了 codelab。

• notebook：https://github.com/GoogleCloudPlatform/training-data-analyst/blob/master/quests/tpu/flowers_resnet.ipynb
• codelab：https://codelabs.developers.google.com/codelabs/tpu-resnet
• Cloud Datalab：https://cloud.google.com/datalab
• Cloud Shell：https://cloud.google.com/shell/

Cloud TPUv2（如上图所示）可以加快最先进的深度学习模型的训练

教程目录

1. 指向 JPEG 数据的 CSV 文件
2. 启用 Cloud TPU 服务账号
3. 复制 ResNet 代码
4. [可选] 在本地尝试数据预处理
5. 运行数据预处理代码
6. 训练模型
7. 部署模型
8. 用模型进行预测

1. 指向 JPEG 数据的 CSV 文件

开始之前，你需要一个装满图像文件和三个逗号分隔值（CSV）文件的文件夹，这些文件提供关于图像的元数据。

首先，你需要一个 CSV 文件，该文件包含你希望用于训练的图像及其标签。CSV 文件的每一行可能如下所示：

gs://cloud-ml-data/img/30a9ae018c_n.jpg,daisy
gs://cloud-ml-data/img/907ed2c7cd_m.jpg,dandelion
gs://cloud-ml-data/img/93a01f48f6.jpg,dandelion
gs://cloud-ml-data/img/81d8c9b0bd_m.jpg,dandelion

你可以根据自己的喜好对图片进行命名，但是它的路径应该是实时的，并且可以在谷歌云存储上访问。标签字符串也可以是你喜欢的任何字符串，但其中不能包含逗号。数据中应该至少包含两类图像，并且训练数据集应该包含足够多的每个类别的示例。因为本文需要从头开始做图像分类工作，我建议每个类别至少有 1000 张图像，总的数据集规模至少为 20,000 张图像。如果你的图像数量较少，可以参考迁移学习教程（它使用相同的数据格式）。

其次，你需要一个和上面一样的 CSV，然后将其用来评估模型。我建议你将 90% 的数据用于训练，而另外 10% 的数据用于评估。确保评估数据集包含每个类别 10% 的图像。

最后，你需要一个包含所有唯一标签的文件，每行一个标签。例如：

daisy
dandelion
roses
sunflowers
tulips

文件中的顺序非常重要。如果最终的模型预测结果为「2」，你就可以得知该图像为玫瑰（第 0 类代表雏菊）。你可以从你用于训练的 CSV 文件中得到类的列表：

gsutil cat gs://cloud-ml-data/img/flower_photos/train_set.csv \ | sed 's/,/ /g' | awk '{print $2}' \ | sort | uniq > /tmp/labels.txt

在上面的代码中，我仅仅从训练 CSV 文件中提取出了第二个字段，并且对其进行排序，在得到的输出结果中寻找到这些值的唯一的集合。通过你最熟悉的过程创建下面三个文件：「train_set.csv」、「eval_set.csv」、「labels.txt」，将他们上传到云存储中，然后你就做好训练模型的准备工作了。

2. 复制 ResNet 代码

让我们从官方 TPU 样本（https://medium.com/r/?url=https%3A%2F%2Fgithub.com%2Ftensorflow%2Ftpu）中复制 ResNet 代码，并制作一个可提交的包。为了做到这一点，你需要从我的 GitHub 代码仓库（https://github.com/GoogleCloudPlatform/training-data-analyst）中复制并运行以下脚本：

git clone https://github.com/GoogleCloudPlatform/training-data-analyst
cd training-data-analyst/quests/tpu/
bash ./copy_resnet_files.sh 1.8

上面的「1.8」指的是 TensorFlow 1.8。我推荐大家使用最新版本的 TensorFlow。

3. 启用 Cloud TPU 服务账号

你需要允许 TPU 服务账号与 ML Engine（机器学习引擎）进行对话。可以使用以下脚本查询服务账号，并且提供访问权限：

bash enable_tpu_mlengine.sh

4. [可选] 在本地尝试进行数据预处理

为了确保我们包的创建工作奏效，你可以尝试运行下面的流程将 JPEG 文件转换为 TensorFlow 记录：

#!/bin/bash
export PYTHONPATH=${PYTHONPATH}:${PWD}/mymodel 
rm -rf /tmp/outpython -m trainer.preprocess \ --train_csv /tmp/input.csv \ --validation_csv /tmp/input.csv \ --labels_file /tmp/labels.txt \ --project_id $PROJECT \ --output_dir /tmp/out --runner=DirectRunner

在这里，「/tmp/input.csv」是你输入的训练文件的一小部分。请检查训练文件和验证文件是否已经被正确创建。

5. 运行预处理代码

运行以下代码将 JPEG 文件转换为 Cloud Dataflow 中的 TFReocord。这将向许多机器分发转换代码，并且自动放缩它的规模：

#!/bin/bash
export PYTHONPATH=${PYTHONPATH}:${PWD}/mymodel
gsutil -m rm -rf gs://${BUCKET}/tpu/resnet/data
python -m trainer.preprocess \ --train_csv gs://cloud-ml-data/img/flower_photos/train_set.csv \ --validation_csv gs://cloud-ml-data/img/flower_photos/eval_set.csv \ --labels_file /tmp/labels.txt \ --project_id $PROJECT \ --output_dir gs://${BUCKET}/tpu/resnet/data

自动放缩 TensorFlow 记录的创建

如果你希望在更新的数据上重新训练你的模型，只需要在新的数据上运行这整套流程，但是请确保将其写入到一个新的输出目录中，以免覆盖之前的输出结果。

6. 训练模型

只需将训练任务提交到 Cloud ML Engine 上，让结果指向你的 Dataflow 作业的输出目录：

#!/bin/bash
TOPDIR=gs://${BUCKET}/tpu/resnet
OUTDIR=${TOPDIR}/trained
REGION=us-central1
JOBNAME=imgclass_$(date -u +%y%m%d_%H%M%S)
echo $OUTDIR $REGION $JOBNAME
gsutil -m rm -rf $OUTDIR # Comment out this line to continue training from the last time
gcloud ml-engine jobs submit training $JOBNAME \
 --region=$REGION \
 --module-name=trainer.resnet_main \
 --package-path=$(pwd)/mymodel/trainer \
 --job-dir=$OUTDIR \
 --staging-bucket=gs://$BUCKET \
 --scale-tier=BASIC_TPU \
 --runtime-version=$TFVERSION \
 -- \
 --data_dir=${TOPDIR}/data \
 --model_dir=${OUTDIR} \
 --resnet_depth=18 \
 --train_batch_size=128 --eval_batch_size=32 --skip_host_call=True \
 --steps_per_eval=250 --train_steps=1000 \
 --num_train_images=3300 --num_eval_images=370 --num_label_classes=5 \
 --export_dir=${OUTDIR}/export

（以上是代码截图）

代码中加粗的行代表了你可能想进行调整的部分：

1. 通过这一行，我们可以在启动训练作业之前删除「OUTDIR」。这会让训练从头重新开始。如果你有新的图像需要训练，并且只希望更新现有的模型，那么不需要删除输出目录。
2. 在这里，我们使用了 ResNet-18，它是最小的 ResNet 模型。你可以选择 ResNet-18、34、50 等模型。（完整列表请参阅「resnet_main.py」：https://medium.com/r/?url=https%3A%2F%2Fgithub.com%2Ftensorflow%2Ftpu%2Fblob%2Fmaster%2Fmodels%2Fofficial%2Fresnet%2Fresnet_main.py）。随着数据集规模的增大，这些数据可以支撑起越来越大的模型的训练：较大的模型在较小的数据集上进行训练存在过拟合的风险。因此随着数据集大小的增加，你可以使用更大的模型。
3. 张量处理单元（TPU）在批处理（batch）规模为 1024 左右时工作效果非常好。而我所拥有的数据集非常小，因此使用较小的批处理规模的原因。
4. 「train_steps」变量控制着你计划用于训练的时间（多少轮迭代）。每次给模型输入数量为「train_batch_size」的图像。要想得到一个大致合理的值，你可以尝试配置你的训练会话（session），这样模型至少能接收到每个图像 10 次。在本文的例子中，我拥有 3,300 张图像，「train_batch_size」为 128，因此，为了模型能接收到每张图像 10 次，我需要（3300*10）/128 步或者大约 250 步。损失曲线（见下一节 TensorBoard 中的示意图）在 250 步时并没有停滞（收敛），所以我将该值增大到 1,000。
5. 「steps_per_eval」变量控制了评估的频率。进行模型评估的计算开销是高昂的，所以你需要试着限制评估的次数。我将训练步设为 1000，每 250 步进行一次评估，因此我将对模型进行 4 次评估。
6. 你需要明确指定训练图像、评估图像以及标签的数量。我使用以下脚本来确定这些数字（通过改变文件名指向你的数据集）：

#!/bin/bash
echo -n "--num_train_images=$(gsutil cat gs://cloud-ml-data/img/flower_photos/train_set.csv | wc -l) "
echo -n "--num_eval_images=$(gsutil cat gs://cloud-ml-data/img/flower_photos/eval_set.csv | wc -l) "
echo "--num_label_classes=$(cat /tmp/labels.txt | wc -l)"

当模型训练完成后（这取决于训练文件批处理规模「train_batch_size」以及训练步「train_step」的数量），模型文件将被导出至谷歌云存储中。

你可以通过 TensorBoard 查看最终的模型的质量（令其指向输出目录）：

没有严重的过拟合现象——损失曲线和评估准确率大致相等

准确率确实太低了，只有 80%。如果使用更多的数据进行训练将有助于准确率提升。

7. 部署模型

你现在可以将模型作为 web 服务部署到 Cloud ML Engine 上（或者你可以自行安装 TensorFlow Serving，并且在其他地方运行模型）：

#!/bin/bash
MODEL_NAME="flowers"
MODEL_VERSION=resnet
MODEL_LOCATION=$(gsutil ls gs://${BUCKET}/tpu/resnet/trained/export/ | tail -1
echo "Deleting and deploying $MODEL_NAME $MODEL_VERSION from $MODEL_LOCATION ... this will take a few minutes"
gcloud ml-engine models create ${MODEL_NAME} --regions $REGION
gcloud ml-engine versions create ${MODEL_VERSION} --model ${MODEL_NAME} --origin ${MODEL_LOCATION} --runtime-version=$TFVERSION

8. 通过模型进行预测

想要使用该模型进行预测，你需要将一个通过 base-64 方式编码的 JPEG 图像文件的内容发送到 web 服务上。下面是创建必要的字典数据结构的 Python 代码片段：

{"image_bytes": {"b64": base64.b64encode(filecontents)}}

将代码封装到可以进行必要身份验证和 HTTP 调用的模版中：

from googleapiclient import discoveryfrom oauth2client.client import GoogleCredentialsimport base64, sys, jsonimport tensorflow as tfwith tf.gfile.FastGFile('gs://cloud-ml-data/img/flower_photos/sunflowers/1022552002_2b93faf9e7_n.jpg', 'r') as ifp: 
 credentials = GoogleCredentials.get_application_default() 
 api = discovery.build('ml', 'v1', credentials=credentials, 

discoveryServiceUrl='*https://storage.googleapis.com/cloud-ml/discovery/ml_v1_discovery.json'* (https://medium.com/r/?url=https%3A%2F%2Fstorage.googleapis.com%2Fcloud-ml%2Fdiscovery%2Fml_v1_discovery.json%27))

 request_data = {'instances': [ {"image_bytes": {"b64": base64.b64encode(ifp.read())}} ]}
parent = 'projects/%s/models/%s/versions/%s' % (PROJECT, 'flowers', 'resnet')
 response = api.projects().predict(body=request_data, name=parent).execute() print "response={0}".format(response)

当我使用这张图片调用该模型时，得到了预期结果（向日葵）：

这是向日葵还是其它的花呢？

原文链接：https://cloud.google.com/blog/big-data/2018/07/how-to-train-a-resnet-image-classifier-from-scratch-on-tpus-on-cloud-ml-engine