选自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 性能最好。

例如：

用分片文档数据集（例如 TFRecord）；用 tf.data API 将训练数据传到 TPU；batch size 较大时（如 batch_size=128）。

在工作流程中加入这些步骤，可以避免常见的 I/O 瓶颈，否则这样的 I/O 瓶颈会妨碍 TPU 发挥全部功能。用户可以访问 Kaggle TPU 官方文档，了解其他一些在 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 倍。

官方教程笔记：

图 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]

训练模型时的价格考量

虽然平等地比较了硬件，但它们在价格上有相当大的差异。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 美元奖金（二等奖）的。

文章链接：

该选哪种硬件呢？

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

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

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

原文链接：