Kaggle竞赛硬件如何选择？不差钱、追求速度，那就上TPU吧

选自towardsdatascience

作者：Paul Mooney

机器之心编译

在每种机器学习算法背后，都是以数千兆赫频率运算的硬件提供支持。你可能已经注意到，在设置 Kaggle Notebooks 计算环境时，有几种处理器可供选择，但究竟哪种才是最适合你的呢？本文分别比较了在训练由 tf.keras 编写的机器学习模型过程中，使用 CPU（英特尔至强处理器）、GPU（英伟达 Tesla P100 显卡）和 TPU（谷歌 TPU v3）的优劣势。希望本文可以为 Kaggle 竞赛准备者理解和选择处理器提供帮助。

图 1：在 Kaggle Notebook 中可以免费使用 CPU、GPU 和 TPU。

如何准备测试？

为了比较 CPU、GPU 和 TPU 完成常见数据科学任务时的性能，我们使用 tf_flowers 数据集训练了一个卷积神经网络，然后用三种不同的后端（GPU 是英伟达 Tesla P100+2GHz、双核的英特至强 CPU，13 GB RAM；TPU 是谷歌 8 核的 TPUv3+2GHz、四核的英特尔至强 CPU，16 GB RAM）将相同的代码运行了三遍。所附的 Kaggle 教程笔记（tutorial notebook）显示，在一些情况下，TPU 性能最好。

例如：

1. 用分片文档数据集（例如 TFRecord）；

2. 用 tf.data API 将训练数据传到 TPU；

3. batch size 较大时（如 batch_size=128）。

在工作流程中加入这些步骤，可以避免常见的 I/O 瓶颈，否则这样的 I/O 瓶颈会妨碍 TPU 发挥全部功能。用户可以访问 Kaggle TPU 官方文档，了解其他一些在 TPU 上运行代码的优化技巧。

官网链接：https://www.kaggle.com/docs/tpu

硬件性能如何

这三种硬件之间最明显的区别在于使用 tf.keras 训练模型的速度。tf.keras 库是最流行的机器学习框架之一，因为 tf.keras 可以更快、更方便地实验新想法。如果写代码的时间短，则就可以拨出更多时间进行计算。如果等代码跑完的时间短，也就可以有更多时间评估新想法了。在机器学习比赛中，tf.keras 和 TPU 是一对很强大的组合！

图 2：tf.keras 概貌。(A)tf.keras 是最流行的实现深度学习模型的框架之一；(B) 在设计深度学习模型时，tf.keras 可以快速进行实验和迭代。

在下图 3 第一个实验中，我们在三种硬件上运行了相同的代码（官方教程笔记中的改进版），并且为了避免 CPU 和 GPU 出现内存不足的问题，batch size 也设置得很小（设为 16）。在这种情况下可以看出，训练 Xception 模型时，TPU 比 CPU 快了约 100 倍，比 GPU 快了约 3.5 倍，这是因为 TPU 处理批大小很大的数据时效率更高。我们也试着将 batch size 增加到 128，结果发现 TPU 速度快了 2 倍，GPU 和 CPU 均出现了内存不足的问题。所以，与之前的实验相比，TPU 训练 Xception 模型的速度比 GPU 快 7 倍。

官方教程笔记：https://www.kaggle.com/mgornergoogle/five-flowers-with-keras-and-xception-on-tpu

图 3：用 CPU、GPU 和 TPU 训练 Xception 模型（12 个 epoch）所需时间对比。

如下图 4 所示，模型训练的加速情况也与模型类别有关，Xception 和 Vgg16 就比 ResNet50 表现更好。在这种边界很大的情况下，模型训练速度是 TPU 唯一超过 GPU 的地方。例如，当预测少量样本时，TPU 比 CPU 快了约 3 倍，但又比 GPU 慢了约 3 倍（在某些情况，比如在 batch size 很大的情况下做预测，TPU 表现出色，但本实验中不存在这样的情况）。

图 4：模型类型和任务类型不同，加速情况也不同。(A) 与 ResNet50 相比，Xception 和 Vgg16 的提速更为明显。(B) 在诸如少量样本预测等特定任务中，GPU 的表现比 TPU 好。

为了使结果更为充分，我们注意到 Yu Emma Wang 等人 [1] 开发了一个叫做 ParaDnn 的严格基准，这一基准可以用来比较不同硬件训练机器学习模型时的性能。利用 Yu Emma Wang 等人 [1] 的方法，可以总结出：用 TPU 代替 GPU 时，参数化模型性能可以提升 1 到 10 倍，使用模型的性能也可以提升 3 到 6.8 倍（下图 5）。分片数据、很大的数据批和大模型结合在一起时，TPU 是最佳选择。

图 5：Wang[1] 论文中的重要发现。当 batch size 和 CNN 模型都较大时，TPU 的性能最好。你可以在 Kumar[2] 和 Jouppi[3] 的论文中，找到其他基准实验。

[1] https://arxiv.org/abs/1907.10701

训练模型时的价格考量

虽然平等地比较了硬件，但它们在价格上有相当大的差异。TPU（谷歌 TPU v3 每小时 8.00 美元，GCP 上可以按需选择的 TPU v2 每小时 4.50 美元）比 GPU（英伟达 Tesla P100 每小时 1.46 美元）贵了大概五倍。虽然如此，如果你想优化成本，那还是应该选择 TPU，因为在训练相同模型的情况下，TPU 的训练速度至少快了 5 倍。

当数据以分片格式储存在 GCS bucket，然后以大 batch size 传递到 TPU 时，模型训练会提速约 5 倍，因此建议熟悉 tf.data API 的用户使用 TPU。

有些机器学习实践者优先考虑模型训练时间，而不是模型训练成本。因此，对于想尽快完成模型训练的人来说，TPU 是最佳选择。在训练模型上花更少的时间，就可以花更多时间考虑新想法。但请不要一味听信本文内容——你可以通过 Kaggle Notebooks 在 CPU、GPU 和 TPU 上免费运行代码并评估结果。Kaggle 用户们已经成功在 TPU 上运行了文本数据，并从中获得了许多乐趣。

下面这篇文章的作者描述了是如何用 TPU 训练 BERT 模型，并在最近的 Kaggle 竞赛中赢得了 8000 美元奖金（二等奖）的。

文章链接：https://www.kaggle.com/c/tensorflow2-question-answering/discussion/127333

该选哪种硬件呢？

总之，如果有通用性和大内存需求的话，我们建议你使用 CPU。当你想加快各种数据科学流程时，GPU 是很好的替代品。如果想要尽可能快地完成模型训练的话，TPU 则是最佳选择。

针对所用硬件来优化代码，可以得到更好的结果。我们认为，对分别针对 GPU 代码和 TPU 代码的运行时（runtime）进行比较也很有意思。例如，用像 RAPIDS.ai 这样的 GPU 加速库训练梯度提升模型，再用像 tf.keras 这样的 TPU 加速库训练深度学习模型，比较二者的训练时间，这也是很有意思的。

那么，训练准确的机器学习模型最少需要多长时间？一天内可以评价多少不同的想法？结合 tf.keras，TPU 让机器学习从业人员写代码的时间变短了，等代码跑完的时间也变短了，这样就可以留出更多时间评估新想法，并在 Kaggle 大赛中提升自己作品的性能。

原文链接：https://towardsdatascience.com/when-to-use-cpus-vs-gpus-vs-tpus-in-a-kaggle-competition-9af708a8c3eb

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

✄------------------------------------------------

加入机器之心（全职记者 / 实习生）：hr@jiqizhixin.com

投稿或寻求报道：content@jiqizhixin.com

广告 & 商务合作：bd@jiqizhixin.com

关注公众号：拾黑（shiheibook）了解更多

[广告]赞助链接：

四季很好，只要有你，文娱排行榜：https://www.yaopaiming.com/
让资讯触达的更精准有趣：https://www.0xu.cn/