如何在colab用pytorch_lightning白嫖TPU
首先
我们得熟悉pytorch_lightning以及colab的操作。在使用TPU之前,先挂载好我们的云端磁盘以及配置好pytorch需要的TPU环境。
1
2
3
4
5
6
7
8
9
from
drive.mount('/content/gdrive'
import
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,一个小栗子。
如何在colab用pytorch_lightning白嫖TPU
作者 CheaSim
发布于 2021-01-07
更新于 2021-03-08
许可协议