在Chainer中使用GPU

Posted by 徐志平 on December 14, 2017

在Chainer中使用GPU

在本节中,您将了解以下内容:

Chainer与CuPy的关系

  • CuPy的基础知识
  • Chainer的单GPU使用
  • 模型并行计算的多GPU使用
  • 数据并行计算的多GPU使用

阅读本节后,您将能够:

  • 在支持CUDA的GPU上使用Chainer
  • 在Chainer中编写模型并行计算
  • 在Chainer中编写数据并行计算

Chainer与CuPy的关系

从v2.0.0开始,CuPy变成了一个独立的软件包和仓库。即使您的环境中安装了CUDA,也必须单独安装CuPy才能使用GPU。请参阅启用CUDA / cuDNN支持来设置CUDA支持。

Chainer使用CuPy作为GPU计算的后端。特别的,cupy.ndarray类是Chainer的GPU数组实现。 CuPy支持与Numpy具有兼容接口的功能的一个子集。它使我们能够为CPU和GPU编写一个通用代码。它还支持像PyCUDA一样的用户定义的内核生成,这使我们能够编写专用于GPU的快速实现。

chainer.cuda模块从CuPy中导入许多重要的符号。例如,在Chainer代码中,cupy命名空间被称为cuda.cupy。请注意,即使未安装CUDA,也可以导入chainer.cuda模块。

Chainer使用内存池分配GPU内存。如前面部分所示,Chainer在学习和评估的迭代过程中构造并销毁了许多数组。由于CUDA中的内存分配和释放(即cudaMalloc和cudaFree函数)使CPU和GPU计算同步,这种方式会损害了性能,所以它不太适合CUDA架构。为了避免计算过程中的内存分配和释放,Chainer使用CuPy的内存池作为标准的内存分配器。Chainer将CuPy的默认分配器更改为内存池,因此用户可以直接使用CuPy的功能而不需要处理内存分配器。

cupy.ndarray基础

import cupy
import numpy as np
from chainer import cuda 

有关cupy.ndarray的基本用法,请参阅CuPy的文档。

CuPy是一个GPU数组后端,实现了NumPy接口的一个子集。 cup.ndarray类是其核心,这是numpy.ndarray兼容的GPU替代品。 CuPy在cup.ndarray对象上实现了许多功能。请参阅NumPy API支持的子集的参考。了解NumPy可能有助于利用CuPy的大部分功能。请参阅NumPy文档以了解它。

umpy.ndarray中cupy.ndarray的主要区别在于内容被分配在设备内存上。分配在默认情况下在当前设备上进行。当前设备可以通过cupy.cuda进行更改。设备对象如下:

with cupy.cuda.Device(1):
    x_on_gpu1 = cupy.array([1, 2, 3, 4, 5])

CuPy的大部分操作都是在当前设备上完成的。请注意,它会导致在非当前设备上处理数组时出错。

Chainer提供了一些方便的功能来自动切换和选择设备。例如,chainer.cuda.to_gpu()函数将numpy.ndarray对象复制到指定的设备:

x_cpu = np.ones((5, 4, 3), dtype=np.float32)
x_gpu = cuda.to_gpu(x_cpu, device=1)

它相当于使用CuPy的以下代码:

x_cpu = np.ones((5, 4, 3), dtype=np.float32)
with cupy.cuda.Device(1):
    x_gpu = cupy.array(x_cpu)

移动设备数组到主机可以通过chainer.cuda.to_cpu()完成,如下所示:

x_cpu = cuda.to_cpu(x_gpu)

它相当于使用CuPy的以下代码:

with x_gpu.device:
    x_cpu = x_gpu.get()

这些代码中的with语句需要选择适当的CUDA设备。如果用户只使用一个设备,则不需要这些设备切换。 chainer.cuda.to_cpu()chainer.cuda.to_gpu()函数会自动正确切换当前设备。

cuda.get_device_from_id(1).use()
x_gpu1 = cupy.empty((4, 3), dtype='f')  # 'f' indicates float32

with cuda.get_device_from_id(1):
    x_gpu1 = cupy.empty((4, 3), dtype='f')

with cuda.get_device_from_array(x_gpu1):
    y_gpu1 = x_gpu + 1

由于它接受NumPy数组,我们可以编写一个函数来接受NumPy和CuPy数组的正确设备切换:

def add1(x):
    with cuda.get_device_from_array(x):
        return x + 1

CuPy与NumPy的兼容性使我们能够编写CPU / GPU通用代码。 chainer.cuda.get_array_module()函数可以使它变得容易。这个函数根据参数返回numpy或cupy模块。一个CPU / GPU通用函数的定义如下:

# Stable implementation of log(1 + exp(x))
def softplus(x):
    xp = cuda.get_array_module(x)
    return xp.maximum(0, x) + xp.log1p(xp.exp(-abs(x)))

在单个GPU上运行神经网络

单GPU使用非常简单。你需要做的是事先将链接和输入数组传输到GPU。在本小节中,代码基于本教程中的第一个MNIST示例。

import numpy as np
import chainer
from chainer import cuda, Function, gradient_check, report, training, utils, Variable
from chainer import datasets, iterators, optimizers, serializers
from chainer import Link, Chain, ChainList
import chainer.functions as F
import chainer.links as L
from chainer.training import extensions
train, test = datasets.get_mnist()
train
<chainer.datasets.tuple_dataset.TupleDataset at 0x13a31ec18>
train_iter = iterators.SerialIterator(train, batch_size=100, shuffle=True)
test_iter = iterators.SerialIterator(test, batch_size=100, repeat=False, shuffle=False)
class MLP(Chain):
    def __init__(self, n_units, n_out):
        super(MLP, self).__init__()
        with self.init_scope():
            # the size of the inputs to each layer will be inferred
            self.l1 = L.Linear(None, n_units)  # n_in -> n_units
            self.l2 = L.Linear(None, n_units)  # n_units -> n_units
            self.l3 = L.Linear(None, n_out)    # n_units -> n_out

    def __call__(self, x):
        h1 = F.relu(self.l1(x))
        h2 = F.relu(self.l2(h1))
        y = self.l3(h2)
        return y
class Classifier(Chain):
    def __init__(self, predictor):
        super(Classifier, self).__init__()
        with self.init_scope():
            self.predictor = predictor

    def __call__(self, x, t):
        y = self.predictor(x)
        loss = F.softmax_cross_entropy(y, t)
        accuracy = F.accuracy(y, t)
        report({'loss': loss, 'accuracy': accuracy}, self)
        return loss
model = L.Classifier(MLP(100, 10))  # the input size, 784, is inferred
optimizer = optimizers.SGD()
optimizer.setup(model)

Link对象可以使用to_gpu()方法传输到指定的GPU。

这一次,我们可以配置输入,隐藏和输出单元的数量。 to_gpu()方法也接受类似于model.to_gpu(0)的设备ID。在这种情况下,Link对象被传送到适当的GPU设备。当前设备默认使用。

如果我们使用chainer.training.Trainer,我们要做的只是让updater知道设备ID来发送每个小批量。

updater = training.StandardUpdater(train_iter, optimizer, device=0)
trainer = training.Trainer(updater, (20, 'epoch'), out='result')

我们还必须为评估扩展指定设备ID。

trainer.extend(extensions.Evaluator(test_iter, model, device=0))

当我们写下训练循环时,我们必须手动将每个小批量传输到GPU:

model.to_gpu(device=0)

<chainer.links.model.classifier.Classifier at 0x13a325fd0>
from chainer.datasets import mnist

train,test = mnist.get_mnist()
batch = train_iter.next()
from chainer.dataset import convert
x_array, t_array = convert.concat_examples(batch, device=0)
datasize = len(x_array)
chainer.cuda.get_device_from_id(0).use()
for epoch in range(20):
    print('epoch %d' % epoch)
    indexes = np.random.permutation(datasize)
    for i in range(0, datasize, batchsize):
        x = Variable(cuda.to_gpu(x_array[indexes[i : i + batchsize]],device=0))
        t = Variable(cuda.to_gpu(t_array[indexes[i : i + batchsize]],device=0))
        optimizer.update(model, x, t)
epoch 0
epoch 1
epoch 2
epoch 3
epoch 4
epoch 5
epoch 6
epoch 7
epoch 8
epoch 9
epoch 10
epoch 11
epoch 12
epoch 13
epoch 14
epoch 15
epoch 16
epoch 17
epoch 18
epoch 19

多GPU上的模型并行计算

机器学习的并行化大致分为“模型 - 并行”和“数据并行”两种。模型平行意味着模型内部计算的并行化。相反,数据并行意味着使用数据分片的并行化。在本小节中,我们将展示如何在Chainer中的多个GPU上使用模型并行方法。

回想一下MNIST的例子。现在假设我们想要修改这个例子,将网络扩展到6层,每层使用2000个单元,并且使用两个GPU。为了使多GPU计算效率更高,我们只让两个GPU在第三层和第六层进行通信。整体架构如下图所示:

(GPU0) 输入  --+--> l1 --> l2 --> l3 --+--> l4 --> l5 --> l6 --+--> 输出
               |                       |                       |
(GPU1)         +--> l1 --> l2 --> l3 --+--> l4 --> l5 --> l6 --+

如下图,我们可以使用上面的MLP链:

(GPU0) 输入   --+--> mlp1 --+--> mlp2 --+--> 输出
               |           |           |
(GPU1)         +--> mlp1 --+--> mlp2 --+

我们来写一个整个网络的连接。

class ParallelMLP(Chain):
    def __init__(self):
        super(ParallelMLP, self).__init__(
            # the input size, 784, is inferred
            mlp1_gpu0=MLP(1000, 2000).to_gpu(0),
            mlp1_gpu1=MLP(1000, 2000).to_gpu(1),

            # the input size, 2000, is inferred
            mlp2_gpu0=MLP(1000, 10).to_gpu(0),
            mlp2_gpu1=MLP(1000, 10).to_gpu(1),
        )

    def __call__(self, x):
        # assume x is on GPU 0
        z0 = self.mlp1_gpu0(x)
        z1 = self.mlp1_gpu1(F.copy(x, 1))

        # sync
        h0 = F.relu(z0 + F.copy(z1, 0))
        h1 = F.relu(z1 + F.copy(z0, 1))

        y0 = self.mlp2_gpu0(h0)
        y1 = self.mlp2_gpu1(h1)

        # sync
        y = y0 + F.copy(y1, 0)
        return y  # output is on GPU0

回想一下,Link.to_gpu()方法返回链接本身。 copy()函数将输入变量复制到指定的GPU设备,并在设备上返回一个新变量。该副本支持反向传输,只是将输出渐变传输到输入设备。

以上代码不是在CPU上并行化,而是在GPU上并行化。这是因为上述代码中的所有功能都与主机CPU异步运行。

examples/mnist/train_mnist_model_parallel.py中可以找到几乎相同的示例代码。

带有Trainer的在多GPU上的数据并行计算

数据并行计算是并行在线处理的另一种策略。在神经网络的情况下,这意味着不同的设备对输入数据的不同子集进行计算。在本小节中,我们将回顾在两个GPU上实现数据并行学习的方法。

再假设我们的任务是MNIST的例子。这次我们要直接并行三层网络。数据并行化最简单的形式是并行化一组不同的数据的梯度计算。首先,定义一个模型和优化器实例:

model = L.Classifier(MLP(1000, 10))  # the input size, 784, is inferred
optimizer = optimizers.SGD()
optimizer.setup(model)

回想一下,MLP连接实现了多层感知器,分类器连接包装它以提供分类器接口。在前面的例子中我们使用了StandardUpdater。为了启用多个GPU的数据并行计算,我们只需要用ParallelUpdater替换它。

updater = training.ParallelUpdater(train_iter, optimizer,
                                   devices={'main': 0, 'second': 1})

devices选项指定在数据并行学习中使用哪些设备。名称为“main”的设备被用作主设备。原始模型发送到此设备,所以优化运行在主设备上。在上面的例子中,模型也被克隆并发送到GPU 1。每个小批量的一半被馈送到这个克隆的模型。每次反向计算后,梯度累积到主设备中,参数更新运行,更新后的参数再次发送给GPU 1。

另请参阅examples/mnist/train_mnist_data_parallel.py中的示例代码。

在没有Trainer的在多GPU上的数据并行计算

我们在这里介绍一种在没有Trainer的帮助下编写数据并行计算的方法。大多数用户可以跳过这一节。如果您对如何自己编写数据并行计算感兴趣,本节应该提供信息。例如,定制ParallelUpdater类也是有帮助的。

我们再次从MNIST的例子开始。此时,我们使用后缀_0和_1来区分每个设备上的对象。首先,我们定义一个模型。

model_0 = L.Classifier(MLP(1000, 10))  # the input size, 784, is inferred

我们想在不同的GPU上制作这个实例的两个副本。 取而Link.to_gpu()代之的方法是Link.copy()用它来创建一个副本。

import copy
model_1 = copy.deepcopy(model_0)
model_0.to_gpu(0)
model_1.to_gpu(1)
<chainer.links.model.classifier.Classifier at 0x13a325cc0>

Link.copy()方法将连接复制到另一个实例中。它只是复制连接层次结构,并不复制它保存的数组。

然后,建立一个优化器:

optimizer = optimizers.SGD()
optimizer.setup(model_0)

在这里,我们使用模型的第一个副本作为主模型。在更新之前,model_1的梯度必须聚合到model_0的梯度。

然后,我们可以编写一个数据并行学习循环如下:

batchsize = 100
datasize = len(x_train)
for epoch in range(20):
    print('epoch %d' % epoch)
    indexes = np.random.permutation(datasize)
    for i in range(0, datasize, batchsize):
        x_batch = x_train[indexes[i : i + batchsize]]
        y_batch = y_train[indexes[i : i + batchsize]]

        x0 = Variable(cuda.to_gpu(x_batch[:batchsize//2], 0))
        t0 = Variable(cuda.to_gpu(y_batch[:batchsize//2], 0))
        x1 = Variable(cuda.to_gpu(x_batch[batchsize//2:], 1))
        t1 = Variable(cuda.to_gpu(y_batch[batchsize//2:], 1))

        loss_0 = model_0(x0, t0)
        loss_1 = model_1(x1, t1)

        model_0.cleargrads()
        model_1.cleargrads()

        loss_0.backward()
        loss_1.backward()

        model_0.addgrads(model_1)
        optimizer.update()

        model_1.copyparams(model_0)

不要忘记清除两个模型副本的梯度!一半的小批量被转发到GPU0,另一半转到GPU1。然后,通过Link.addgrads()方法累积梯度。这种方法将给定连接的梯度添加到自己的连接。在梯度准备好之后,我们可以用通常的方式更新优化器。请注意,更新仅修改了model_0的参数。所以我们必须使用Link.copyparams()方法手动将它们复制到model_1。

如果在一个模型中使用的批量大小保持不变,则当我们通过chainer.Link.addgrads()从所有模型聚合梯度时,梯度的比例大致与模型的数量成正比。所以您需要相应地调整优化器的批量大小和/或学习速率。

现在你可以在Chainer中使用GPU。示例目录中的所有示例都支持GPU计算,所以如果您想了解有关使用GPU的更多实践,请参阅它们。在下一节中,我们将介绍如何在Variable对象上定义一个可微分(即可反向)的函数。我们还将展示如何使用Chainer的CUDA实用程序编写简单的(元素级别的)CUDA内核。