TPU芯片介绍

Google定制的打机器学习专用晶片称之为TPU（Tensor Processing Unit），Google在其自家称，由于TPU专为机器学习所运行，得以较传统CPU、 GPU 降低精度，在计算所需的电晶体数量上，自然可以减少，也因此，可从电晶体中挤出更多效能，每秒执行更复杂、强大的机器学习模组，并加速模组的运用，使得使用者更快得到答案，Google最早是计划用FPGA的，但是财大气粗，考虑到自己的特殊应用，就招了很多牛人来做专用芯片TPU。

TPUs已经针对TensorFlow进行了优化，并且主要用于TensorFlow。但是Kaggle和谷歌在它的一些比赛中分发了免费的TPU时间，并且一个人不会简单地改变他最喜欢的框架，所以这是一个关于我在GCP上用TPU训练PyTorch模型的经验的备忘录(大部分是成功的)。

PyTorch/XLA是允许这样做的项目。它仍在积极的开发中，问题得到了解决。希望在不久的将来，运行它的体验会更加顺畅，一些bug会得到修复，最佳实践也会得到更好的交流。

https://github.com/pytorch/xla

设置

这里有两种方法可以获得TPU的使用权

GCP计算引擎虚拟机与预构建的PyTorch/XLA映像并按照PyTorch/XLA github页面上的“使用预构建的计算VM映像”部分进行设置。

或者使用最简单的方法，使用google的colab笔记本可以获得免费的tpu使用。

针对一kaggle的比赛您可以在虚拟机上使用以下代码复制Kaggle API令牌并使用它下载竞争数据。还可以使用gsutil cp将文件复制回GS bucket。

 gcloud auth login
 gsutil cp gs://bucket-name/kaggle-keys/kaggle.json ~/.kaggle
 chmod 600 ~/.kaggle/kaggle.json
 kaggle competitions download -c recursion-cellular-image-classification

除了谷歌存储之外，我还使用github存储库将数据和代码从我的本地机器传输到GCP虚拟机，然后再返回。

注意，在TPU节点上也有运行的软件版本。它必须匹配您在VM上使用的conda环境。由于PyTorch/XLA目前正在积极开发中，我使用最新的TPU版本:

使用TPU训练

让我们看看代码。PyTorch/XLA有自己的多核运行方式，由于TPUs是多核的，您希望利用它。但在你这样做之前，你可能想要把你的模型中的device = ' cuda '替换为

 import torch_xla_py.xla_model as xm...
 device = xm.xla_device()...
 xm.optimizer_step(optimizer)