如何在colab用pytorch_lightning白嫖TPU

首先

我们得熟悉pytorch_lightning以及colab的操作。在使用TPU之前，先挂载好我们的云端磁盘以及配置好pytorch需要的TPU环境。

from google.colab import drive
drive.mount('/content/gdrive')
import os
os.chdir("gdrive/MyDrive/path_to_your_file")
! pip install pytorch_lightning

# 试图使用TPU
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py  --apt-packages libomp5 libopenblas-dev

安装好pytorch_lightning以及配置好环境之后，我们就可以开始在pytorch_lightning框架下跟写GPU代码一样写TPU代码。下面请看VCR，一个小栗子。