April 23, 2020

Colab、PyTorch 和 TPU

Colab、PyTorch 和 TPU

机器学习最重要的是获取更多的数据和算力,怎样获取更多数据说多了都是麻烦,目前精力应该更多的集中在获取算力上,目前国内百度的 aistudio 算比较慷慨,有 32g 的 V100 提供,但各种时间限制,只能跑跑 demo,真有项目需求的话只能选择多和他们自己的分布式环境或者他们的 easydl,说实话易用性还有很大提高空间,两种方式都一样。

墙外的无论是 kaggle 还是 colab / GCP 现在都是一家的,Google 的,他家的 TPU 虽然不卖硬件有断供之嫌,但目前提供出来的公用资源确实不错,值得一试,支持最好的当然是 GCP + TPU,价格小贵,个人或者小团体的话每月 10 刀的 Colab Pro 也可以考虑,截至此时此刻(2020.04.23),Google 在 Colab 提供的 TPU 资源是 TPU V2 设备,每个设备下 4 个 TPU V2 模块,每个模块有两个 core,具体到代码的话就是可选用 1 个 或者 8 个 core,其他的选项不支持。粗步算下来相当于 5 个 v100?(粗略估计,可能不对待查待确认)。熟悉了如何在 GCP/Colab 上获取 TPU 的算力之后,可以很方便的通过 PyTorch 或者 Keras/Tensorflow 去 scale 到更多的 TPU 设施上只要你有足够的钱去买服务就好,从目前各家的硬件上看,TPU 算是性价比比较高的。


接下来看下如何在 Colab 上开启 PyTorch 对 TPU 的支持,PyTorch 是通过 XLA 编译器来支持 TPU 的,详情自查。这里只需要知道把下面这个代码放到一个代码块里面执行以下就打开了支持就好:

import os
assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator

VERSION = "20200325"  #@param ["1.5" , "20200325", "nightly"]
!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 --version $VERSIO

通过 torch_xla 包获取 TPU 硬件对 PyTorch 的支持:

# imports pytorch
import torch

# imports the torch_xla package
import torch_xla
import torch_xla.core.xla_model as xm

通过 xladevice() 方法拿到 TPU Core,默认是一个 core,该方法通过 n 和 devkind 两个参数来指定 8 个 core 中你要的 core,默认是第一个(xla:1),没错 index 是从 1 开始的。如下图中,如果你不指定参数,xla_device() 默认会返回第一个 core。

second_dev = xm.xla_device(n=2, devkind='TPU')
t2 = torch.zeros(3, 3, device = second_dev)
print(t2)

接下来在第一个 core 上我们做几个张量的常用操作:

  • 首先获得这个 core:
dev = xm.xla_device()
  • 在这个 core 上做一个简单的求导:
dev = xm.xla_device()
x = torch.rand(5,5, requires_grad=True, device=dev)
y = torch.rand(5,5, requires_grad=True, device=dev)

z = x**2 + y**3
z.backward(torch.ones_like(x))

print(x)
print(x.grad)
print(x.grad/x)
  • 再来个一维卷积:
filters = torch.randn(33, 16, 3, device = dev)
inputs = torch.randn(20, 16, 50, device = dev)
torch.nn.functional.conv1d(inputs, filters
  • 在 cpu 和各个 core 之间拷贝张量玩,注意这里的 3 个张量虽然长得像,但却分别是在内存/tpu1/tpu2 上的 3 个不同向量:
  • 线性变换:
  • 再看下这个简单的 PyTorch 卷积网络,你会发现如果只用一个 core 的话,tpu 跟 gpu/cpu 没啥区别,而且花在初始化环境上的时间还比他们要多:
import torch.nn as nn
import torch.nn.functional as F

# Simple example network from 
# https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html#sphx-glr-beginner-blitz-neural-networks-tutorial-py
class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 6*6 from image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square you can only specify a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


# Places network on the default TPU core
net = Net().to(dev)

# Creates random input on the default TPU core
input = torch.randn(1, 1, 32, 32, device=dev)

# Runs network
out = net(input)
print(out)

BTW,有一个有意思的操作是这样的,xm.xrt_worldsize() 对应的是 TPU core 的数量,想下面这种处理之后是不是有一种即享受了加大 lr 带来更快训练速度,又不用担心 lr 太大而迟迟不肯收敛的快感?(值得注意的一点是,如果你直接 import 了 xm,然后执行 xrt_world_size() 返回是 1  而不是 8)

  # Scale learning rate to world size
  lr = FLAGS['learning_rate'] * xm.xrt_world_size()