weblog

Appearance

异步计算 ​

一、什么是异步计算 ​

我们现在用的计算机都是高度并行的系统，有多个 CPU 核、多个 GPU，每个设备都能同时处理很多事情。但 Python 本身是单线程的，不擅长并行和异步代码，所以深度学习框架（比如 PyTorch、MXNet、飞桨）都用了异步计算的方式来提高性能。

简单来说，异步计算就是：当我们调用一个 GPU 操作的时候，这个操作不会立刻执行，而是先排队到 GPU 的任务队列里，Python 代码会继续往下跑，不用等这个操作完成。这样 CPU 和 GPU 可以同时干活，提高整体的计算效率。

比如我们用 PyTorch 做矩阵乘法，当我们写b = torch.mm(a, a)的时候，Python 只是把这个任务交给 GPU 的后端队列，然后就继续执行下一行代码了，不用等 GPU 算完这个乘法。

二、异步计算的好处：为什么要这么做 ​

我们先看一个例子，对比 NumPy 和 PyTorch 的矩阵乘法速度：

python
import numpy as np
import torch
from d2l import torch as d2l

# 用GPU计算
device = d2l.try_gpu()

# 测试NumPy（CPU计算）
with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = np.random.normal(size=(1000, 1000))
        b = np.dot(a, a)

# 测试PyTorch（GPU异步计算）
with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)

运行结果大概是：

Plain
numpy: 1.0704 sec
torch: 0.0013 sec

PyTorch 快了好几个数量级，除了 GPU 本身比 CPU 快，还有一个重要原因是 PyTorch 用了异步计算：Python 把任务交给 GPU 后，就继续干别的了，不用等 GPU 算完，所以看起来耗时特别短。

如果我们强制让 Python 等 GPU 算完（同步），耗时就会变长：

python
with d2l.Benchmark():
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
    # 强制等待GPU计算完成
    torch.cuda.synchronize(device)

运行结果大概是Done: 0.0049 sec，比异步的时候长了一点，但还是比 NumPy 快很多，因为 GPU 本身的计算速度快。

三、异步计算的原理：前端和后端 ​

深度学习框架分为两部分：

  1. 前端：就是我们写代码的地方，比如 Python，负责发出计算命令。

  2. 后端：是框架的 C++ 实现，负责实际执行计算，管理自己的线程，不断从队列里拿任务执行。

前端把计算任务发给后端后，就继续执行自己的代码，不用等后端完成。后端会跟踪任务之间的依赖关系，比如如果z = x * y + 2，那么后端会先算x*y，再算+2，不会把顺序搞乱。

四、阻塞操作：什么时候会强制等待 ​

有时候我们需要拿到计算结果，这时候就必须等后端算完，这种操作叫阻塞操作，比如：

  1. 打印变量：print(z)，因为要打印结果，必须等 z 算完。

  2. 把张量转换成 NumPy：z.numpy()，因为 NumPy 没有异步的概念，必须拿到实际的数值。

  3. 把张量转换成标量：z.item()，同样需要拿到实际的数值。

  4. 显式等待：torch.cuda.synchronize(device)或者z.wait_to_read()（MXNet 里的方法）。

比如我们测试一下阻塞操作的耗时：

python
# 测试转换成NumPy
with d2l.Benchmark('numpy conversion'):
    b = torch.mm(a, a)
    b.numpy()

# 测试转换成标量
with d2l.Benchmark('scalar conversion'):
    b = torch.mm(a, a)
    b.sum().item()

这些操作都会让 Python 等待 GPU 计算完成，所以耗时会比纯异步的时候长。

五、异步计算的性能提升 ​

异步计算可以让前端和后端并行工作，减少等待时间。比如我们做 10000 次加法操作，对比同步和异步的耗时：

python
# 同步执行：每次都等计算完成
with d2l.Benchmark('synchronous'):
    for _ in range(10000):
        y = x + 1
        # 每次都等y算完
        torch.cuda.synchronize(device)  # 强制等待

# 异步执行：最后再等所有计算完成
with d2l.Benchmark('asynchronous'):
    for _ in range(10000):
        y = x + 1
    # 最后等所有计算完成
    torch.cuda.synchronize(device)

运行结果大概是：

Plain
synchronous: 3.5791 sec
asynchronous: 0.7830 sec

异步的耗时只有同步的 1/4 左右，因为同步的时候，Python 每次都要等 GPU 算完才能继续下一次循环，而异步的时候，Python 可以一次性把所有任务都发给 GPU，然后 GPU 自己慢慢算，Python 不用等。

假设每个任务的三个阶段耗时分别是 t1（前端发任务）、t2（后端计算）、t3（后端返回结果）：

  • 同步的总时间是：10000*(t1 + t2 + t3)

  • 异步的总时间是：t1 + 10000t2 + t3（只要 10000t2 > 9999*t1，异步就更快）

六、注意事项 ​

  1. 不要过度异步：如果把太多任务都放到队列里，会占用太多内存，导致内存不够用。建议每个小批量训练完后，做一次同步，让前端和后端的进度保持一致。

  2. 避免频繁的阻塞操作：如果频繁把张量转换成 NumPy 或者打印结果，会破坏异步的性能，因为每次都要等 GPU 算完。

  3. 注意任务依赖：后端会自动跟踪任务的依赖关系，所以不用担心顺序错了，比如z = x*y + 2，后端会先算x*y，再算+2，不会乱序。

七、小结 ​

  1. 异步计算让 Python 前端和框架后端并行工作，提高计算效率，尤其是在 GPU 计算的时候。

  2. 阻塞操作会强制等待计算完成，比如打印、转换 NumPy、显式同步，要合理使用。

  3. 不要过度异步，避免内存占用过多，建议每个小批量同步一次。

（注：文档部分内容可能由 AI 生成）

京ICP备2024093538号-1

自用笔记本