异步计算

MXNet 使用异步计算来提升计算性能。理解它的工作原理既有助于开发更高效的程序,又有助于在内存资源有限的情况下主动降低计算性能从而减小内存开销。

我们先导入本节中实验需要的包或模块。

In [1]:
from mxnet import autograd, gluon, nd
from mxnet.gluon import loss as gloss, nn
import os
import subprocess
from time import time

MXNet 中的异步计算

广义上,MXNet 包括用户直接用来交互的前端和系统用来执行计算的后端。例如,用户可以使用不同的前端语言编写 MXNet 程序,像 Python、R、Scala 和 C++。无论使用何种前端编程语言,MXNet 程序的执行主要都发生在 C++ 实现的后端。换句话说,用户写好的前端 MXNet 程序会传给后端执行计算。后端有自己的线程在队列中不断收集任务并执行它们。

MXNet 通过前端线程和后端线程的交互实现异步计算。异步计算指,前端线程无需等待当前指令从后端线程返回结果就继续执行后面的指令。为了便于解释,假设 Python 前端线程调用以下四条指令。

In [2]:
a = nd.ones((1, 2))
b = nd.ones((1, 2))
c = a * b + 2
c
Out[2]:

[[ 3.  3.]]
<NDArray 1x2 @cpu(0)>

在异步计算中,Python 前端线程执行前三条语句的时候,仅仅是把任务放进后端的队列里就返回了。当最后一条语句需要打印计算结果时,Python 前端线程会等待 C++ 后端线程把 c 的结果计算完。此设计的一个好处是,这里的 Python 前端线程不需要做实际计算。因此,无论 Python 的性能如何,它对整个程序性能的影响很小。只要 C++ 后端足够高效,那么不管前端语言性能如何,MXNet 都可以提供一致的高性能。

下面的例子通过计时来展示异步计算的效果。可以看到,当 y = nd.dot(x, x) 返回的时候并没有等待它真正被计算完。

In [3]:
start = time()
x = nd.random.uniform(shape=(2000, 2000))
y = nd.dot(x, x)
print('workloads are queued: %f sec' % (time() - start))
print(y)
print('workloads are completed: %f sec' % (time() - start))
workloads are queued: 0.000511 sec

[[ 501.15838623  508.29724121  495.65237427 ...,  492.8470459   492.69091797
   490.0480957 ]
 [ 508.81057739  507.18218994  495.17428589 ...,  503.10525513
   497.29315186  493.6791687 ]
 [ 489.565979    499.47015381  490.17721558 ...,  490.99945068
   488.05007935  483.2883606 ]
 ...,
 [ 484.00189209  495.71789551  479.92141724 ...,  493.69952393
   478.89193726  487.20739746]
 [ 499.64932251  507.65093994  497.59381104 ...,  493.0473938   500.74511719
   495.82711792]
 [ 516.01428223  519.17150879  506.35400391 ...,  510.08877563  496.3560791
   495.42523193]]
<NDArray 2000x2000 @cpu(0)>
workloads are completed: 0.156485 sec

的确,除非我们需要打印或者保存计算结果,我们基本无需关心目前结果在内存中是否已经计算好了。只要数据是保存在 NDArray 里并使用 MXNet 提供的运算符,MXNet 将默认使用异步计算来获取高计算性能。

用同步函数让前端等待计算结果

除了前面介绍的 print 外,我们还有其他方法让前端线程等待后端的计算结果完成。我们可以使用 wait_to_read 函数让前端等待某个的 NDArray 的计算结果完成,再执行前端中后面的语句。或者,我们可以用 waitall 函数令前端等待前面所有计算结果完成。后者是性能测试中常用的方法。

下面是使用 wait_to_read 的例子。输出用时包含了 y 的计算时间。

In [4]:
start = time()
y = nd.dot(x, x)
y.wait_to_read()
time() - start
Out[4]:
0.12982511520385742

下面是使用 waitall 的例子。输出用时包含了 yz 的计算时间。

In [5]:
start = time()
y = nd.dot(x, x)
z = nd.dot(x, x)
nd.waitall()
time() - start
Out[5]:
0.2595338821411133

此外,任何将 NDArray 转换成其他不支持异步计算的数据结构的操作都会让前端等待计算结果。例如当我们调用 asnumpyasscalar 函数时:

In [6]:
start = time()
y = nd.dot(x, x)
y.asnumpy()
time() - start
Out[6]:
0.13634824752807617
In [7]:
start = time()
y = nd.dot(x, x)
y.norm().asscalar()
time() - start
Out[7]:
0.16842055320739746

上面介绍的 wait_to_readwaitallasnumpyasscalarprint 函数会触发让前端等待后端计算结果的行为,我们通常把这类函数称作同步函数。

使用异步计算提升计算性能

在下面例子中,我们用 for 循环不断对 y 赋值。当 for 循环内使用同步函数 wait_to_read 时,每次赋值不使用异步计算;当 for 循环外使用同步函数 waitall 时,则使用异步计算。

In [8]:
start = time()
for _ in range(1000):
    y = x + 1
    y.wait_to_read()
print('synchronous: %f sec' % (time() - start))

start = time()
for _ in range(1000):
    y = x + 1
nd.waitall()
print('asynchronous: %f sec' % (time() - start))
synchronous: 1.220285 sec
asynchronous: 0.741739 sec

我们观察到,使用异步计算能提升一定的计算性能。为了解释这个现象,让我们对 Python 前端线程和 C++ 后端线程的交互稍作简化。在每一次循环中,前端和后端的交互大约可以分为三个阶段:

  1. 前端令后端将计算任务 y = x + 1 放进队列;
  2. 后端从队列中获取计算任务并执行真正的计算;
  3. 后端将计算结果返回给前端。

我们将这三个阶段的耗时分别设为 \(t_1, t_2, t_3\)。如果不使用异步计算,执行 1000 次计算的总耗时大约为 \(1000 (t_1+ t_2 + t_3)\);如果使用异步计算,由于每次循环前端都无需等待后端返回计算结果,执行 1000 次计算的总耗时可以降为 \(t_1 + 1000 t_2 + t_3\)(假设 \(1000t_2 > 999t_1\))。

异步计算对内存使用的影响

为了解释异步计算对内存使用的影响,让我们先回忆一下前面章节的内容。

在前面章节中实现的模型训练过程中,我们通常会在每个小批量上评测一下模型,例如模型的损失或者精度。细心的你也许发现了,这类评测常用到同步函数,例如 asscalar 或者 asnumpy。如果去掉这些同步函数,前端会将大量的小批量计算任务在极短的时间内丢给后端,从而可能导致较大的内存开销。当我们在每个小批量上都使用同步函数时,前端在每次迭代时仅会将一个小批量的任务丢给后端执行计算,并通常会减小内存开销。

由于深度学习模型通常比较大,而内存资源通常有限,我们建议大家在训练模型时对每个小批量都使用同步函数,例如用 asscalar 或者 asnumpy 评价模型的表现。类似地,在使用模型预测时,为了减小内存开销,我们也建议大家对每个小批量预测时都使用同步函数,例如直接打印出当前小批量的预测结果。

下面我们来演示异步计算对内存使用的影响。我们先定义一个数据获取函数,它会从被调用时开始计时,并定期打印到目前为止获取数据批量总共耗时。

In [9]:
num_batches = 41
def data_iter():
    start = time()
    batch_size = 1024
    for i in range(num_batches):
        if i % 10 == 0:
            print('batch %d, time %f sec' % (i, time() - start))
        X = nd.random.normal(shape=(batch_size, 512))
        y = nd.ones((batch_size,))
        yield X, y

以下定义多层感知机、优化器和损失函数。

In [10]:
net = nn.Sequential()
net.add(
    nn.Dense(2048, activation='relu'),
    nn.Dense(512, activation='relu'),
    nn.Dense(1),
)
net.initialize()
trainer = gluon.Trainer(net.collect_params(), 'sgd',
                        {'learning_rate':0.005})
loss = gloss.L2Loss()

这里定义辅助函数来监测内存的使用。需要注意的是,这个函数只能在 Linux 或 MacOS 运行。

In [11]:
def get_mem():
    res = subprocess.check_output(['ps', 'u', '-p', str(os.getpid())])
    return int(str(res).split()[15]) / 1e3

现在我们可以做测试了。我们先试运行一次让系统把 net 的参数初始化。相关内容请参见 “模型参数的延后初始化” 一节。

In [12]:
for X, y in data_iter():
    break
loss(y, net(X)).wait_to_read()
batch 0, time 0.000002 sec

对于训练 net 来说,我们可以自然地使用同步函数 asscalar 将每个小批量的损失从 NDArray 格式中取出,并打印每个迭代周期后的模型损失。此时,每个小批量的生成间隔较长,不过内存开销较小。

In [13]:
mem = get_mem()
for epoch in range(1, 3):
    l_sum = 0
    for X, y in data_iter():
        with autograd.record():
            l = loss(y, net(X))
        l_sum += l.mean().asscalar()
        l.backward()
        trainer.step(X.shape[0])
    print('epoch', epoch, ' loss: ', l_sum / num_batches)
nd.waitall()
print('increased memory: %f MB' % (get_mem() - mem))
batch 0, time 0.000002 sec
batch 10, time 1.383534 sec
batch 20, time 3.434668 sec
batch 30, time 4.919590 sec
batch 40, time 6.397513 sec
epoch 1  loss:  0.15615208556
batch 0, time 0.000003 sec
batch 10, time 1.492349 sec
batch 20, time 2.976489 sec
batch 30, time 4.663930 sec
batch 40, time 6.253546 sec
epoch 2  loss:  0.0985021460347
increased memory: 11.004000 MB

如果去掉同步函数,虽然每个小批量的生成间隔较短,训练过程中可能会导致内存开销过大。这是因为默认异步计算下,前端会将所有小批量计算在短时间内全部丢给后端。

In [14]:
mem = get_mem()
for epoch in range(1, 3):
    for X, y in data_iter():
        with autograd.record():
            l = loss(y, net(X))
        l.backward()
        trainer.step(x.shape[0])
nd.waitall()
print('increased memory: %f MB' % (get_mem() - mem))
batch 0, time 0.000002 sec
batch 10, time 0.018991 sec
batch 20, time 0.033431 sec
batch 30, time 0.047716 sec
batch 40, time 0.061752 sec
batch 0, time 0.000003 sec
batch 10, time 0.014162 sec
batch 20, time 0.028402 sec
batch 30, time 0.043037 sec
batch 40, time 0.057031 sec
increased memory: 58.172000 MB

小结

  • MXNet 包括用户直接用来交互的前端和系统用来执行计算的后端。
  • MXNet 能够通过异步计算提升计算性能。
  • 我们建议使用每个小批量训练或预测时至少使用一个同步函数,从而避免在短时间内将过多计算任务丢给后端。

练习

  • 在“使用异步计算提升计算性能”一节中,我们提到使用异步计算可以使执行 1000 次计算的总耗时可以降为 \(t_1 + 1000 t_2 + t_3\)。这里为什么要假设 \(1000t_2 > 999t_1\)

扫码直达 讨论区