《动手学深度学习》
Table Of Contents
《动手学深度学习》
Table Of Contents

自动并行计算

MXNet 后端会自动构建计算图。通过计算图,系统可以知道所有计算的依赖关系,并可以选择将没有依赖关系的多个任务并行执行来获得计算性能的提升。例如“异步计算”一节的第一个例子里依次执行了a = nd.ones((1, 2))b = nd.ones((1, 2))。这两步计算之间并没有依赖关系,因此系统可以选择并行执行它们。

通常,一个运算符会用到所有 CPU 或单个 GPU 上全部的计算资源。例如,dot操作符会用到所有 CPU(即使是一台机器上有多个 CPU 处理器)或单个 GPU 上所有的线程。如果每个操作符的计算量足够大,只在 CPU 上或者单个 GPU 上并行运行多个运算符时,每个运算符的运行只分到 CPU 或单个 GPU 上部分计算资源。即使这些计算可以并行,最终计算性能的提升可能并不明显。本节中探讨的自动并行计算主要关注同时使用 CPU 和 GPU 的并行计算,以及计算和通讯的并行。

首先导入本节中实验所需的包或模块。注意,我们需要至少一个 GPU 才能运行本节实验。

In [1]:
import gluonbook as gb
import mxnet as mx
from mxnet import nd

CPU 和 GPU 的并行计算

我们先介绍 CPU 和 GPU 的并行计算,例如程序中的计算既发生在 CPU 上,又发生在 GPU 上。先定义run函数,令它做 10 次矩阵乘法。

In [2]:
def run(x):
    return [nd.dot(x, x) for _ in range(10)]

接下来,分别在 CPU 和 GPU 上创建 NDArray。

In [3]:
x_cpu = nd.random.uniform(shape=(2000, 2000))
x_gpu = nd.random.uniform(shape=(6000, 6000), ctx=mx.gpu(0))

然后,分别使用它们在 CPU 和 GPU 上运行run函数并打印所需时间。

In [4]:
run(x_cpu)  # 预热开始。
run(x_gpu)
nd.waitall()  # 预热结束。

with gb.Benchmark('Run on CPU.'):
    run(x_cpu)
    nd.waitall()

with gb.Benchmark('Then run on GPU.'):
    run(x_gpu)
    nd.waitall()
Run on CPU. time: 1.2315 sec
Then run on GPU. time: 1.2196 sec

我们去掉run(x_cpu)run(x_gpu)两个计算任务之间的nd.waitall(),希望系统能自动并行这两个任务。

In [5]:
with gb.Benchmark('Run on both CPU and GPU in parallel.'):
    run(x_cpu)
    run(x_gpu)
    nd.waitall()
Run on both CPU and GPU in parallel. time: 1.2305 sec

可以看到,当两个计算任务一起执行时,执行总时间小于它们分开执行的总和。这表示,MXNet 能有效地在 CPU 和 GPU 上自动并行计算。

计算和通讯的并行计算

在同时使用 CPU 和 GPU 的计算中,我们经常需要在 CPU 和 GPU 之间复制数据,造成数据的通讯。在下面例子中,我们在 GPU 上计算,然后将结果复制回 CPU。我们分别打印 GPU 上计算时间和 GPU 到 CPU 的通讯时间。

In [6]:
def copy_to_cpu(x):
    return [y.copyto(mx.cpu()) for y in x]

with gb.Benchmark('Run on GPU.'):
    y = run(x_gpu)
    nd.waitall()

with gb.Benchmark('Then copy to CPU.'):
    copy_to_cpu(y)
    nd.waitall()
Run on GPU. time: 1.2254 sec
Then copy to CPU. time: 0.5254 sec

我们去掉计算和通讯之间的waitall函数,打印这两个任务完成的总时间。

In [7]:
with gb.Benchmark('Run and copy in parallel.'):
    y = run(x_gpu)
    copy_to_cpu(y)
    nd.waitall()
Run and copy in parallel. time: 1.2737 sec

可以看到,执行计算和通讯的总时间小于两者分别执行的耗时之和。需要注意的是,这个计算并通讯的任务不同于本节之前介绍的同时使用 CPU 和 GPU 并行计算的任务。这里的运行和通讯之间有依赖关系:y[i]必须先计算好才能复制到 CPU。所幸的是,在计算y[i]的时候系统可以复制y[i-1],从而减少计算和通讯的总运行时间。

小结

  • MXNet 能够通过自动并行计算提升计算性能,例如 CPU 和 GPU 的并行计算以及计算和通讯的并行。

练习

  • 本节中定义的run函数里做了 10 次运算。它们之间也没有依赖关系。设计实验,看看 MXNet 有没有自动并行执行它们。
  • 设计包含更加复杂的数据依赖的计算任务,通过实验观察 MXNet 能否得到正确结果并提升计算性能。
  • 当运算符的计算量足够小时,仅在 CPU 或单 GPU 上并行计算也可能提升计算性能。设计实验来验证这一点。

扫码直达讨论区

image0