在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（如上图所示）可以加快最先进的深度学习模型的训练

教程目录

指向 JPEG 数据的 CSV 文件

启用 Cloud TPU 服务账号

复制 ResNet 代码

[可选] 在本地尝试数据预处理

运行数据预处理代码

训练模型

部署模型

用模型进行预测

1. 指向 JPEG 数据的 CSV 文件

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

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

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

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

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

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

在上面的代码中，我仅仅从训练 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）中复制并运行以下脚本：

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

3. 启用 Cloud TPU 服务账号

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

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

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

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

5. 运行预处理代码

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

自动放缩 TensorFlow 记录的创建

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

6. 训练模型

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

（以上是代码截图）

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

通过这一行，我们可以在启动训练作业之前删除「OUTDIR」。这会让训练从头重新开始。如果你有新的图像需要训练，并且只希望更新现有的模型，那么不需要删除输出目录。

在这里，我们使用了 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）。随着数据集规模的增大，这些数据可以支撑起越来越大的模型的训练：较大的模型在较小的数据集上进行训练存在过拟合的风险。因此随着数据集大小的增加，你可以使用更大的模型。

张量处理单元（TPU）在批处理（batch）规模为 1024 左右时工作效果非常好。而我所拥有的数据集非常小，因此使用较小的批处理规模的原因。

「train_steps」变量控制着你计划用于训练的时间（多少轮迭代）。每次给模型输入数量为「train_batch_size」的图像。要想得到一个大致合理的值，你可以尝试配置你的训练会话（session），这样模型至少能接收到每个图像 10 次。在本文的例子中，我拥有 3,300 张图像，「train_batch_size」为 128，因此，为了模型能接收到每张图像 10 次，我需要（3300*10）/128 步或者大约 250 步。损失曲线（见下一节 TensorBoard 中的示意图）在 250 步时并没有停滞（收敛），所以我将该值增大到 1,000。

「steps_per_eval」变量控制了评估的频率。进行模型评估的计算开销是高昂的，所以你需要试着限制评估的次数。我将训练步设为 1000，每 250 步进行一次评估，因此我将对模型进行 4 次评估。

你需要明确指定训练图像、评估图像以及标签的数量。我使用以下脚本来确定这些数字（通过改变文件名指向你的数据集）：

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

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

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

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

7. 部署模型

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

8. 通过模型进行预测

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

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

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

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

本文为机器之心编译，转载请联系本公众号获得授权。

------------------------------------------------