【动手学深度学习】第十二章笔记:异步计算、数据并行

为了更好的阅读体验,请点击这里

12.1 编译器和解释器

原书主要关注的是命令式编程(imperative programming)。Python 是一种解释性语言,因此没有编译器给代码优化,代码会跑得很慢。

12.1.1 符号式编程

考虑另一种选择符号式编程(symbolic programming),即代码通常只在完全定义了过程之后才执行计算。这个策略被多个深度学习框架使用,包括 Theano 和 TensorFlow(后者已经获得了命令式编程的扩展)。一般包括以下步骤:

  1. 定义计算流程;
  2. 将流程编译成可执行的程序;
  3. 给定输入,调用编译好的程序执行。

这将允许进行大量的优化。首先,在大多数情况下,我们可以跳过 Python 解释器。从而消除因为多个更快的 GPU 与单个 CPU 上的单个 Python 线程搭配使用时产生的性能瓶颈。其次,编译器可以将代码优化和重写。因为编译器在将其转换为机器指令之前可以看到完整的代码,所以这种优化是可以实现的。例如,只要某个变量不再需要,编译器就可以释放内存(或者从不分配内存),或者将代码转换为一个完全等价的片段。下面,我们将通过模拟命令式编程来进一步了解符号式编程的概念。

def add_():
    return '''
def add(a, b):
    return a + b
'''

def fancy_func_():
    return '''
def fancy_func(a, b, c, d):
    e = add(a, b)
    f = add(c, d)
    g = add(e, f)
    return g
'''

def evoke_():
    return add_() + fancy_func_() + 'print(fancy_func(1, 2, 3, 4))'

prog = evoke_()
print(prog)
y = compile(prog, '', 'exec')
exec(y)
def add(a, b):
    return a + b

def fancy_func(a, b, c, d):
    e = add(a, b)
    f = add(c, d)
    g = add(e, f)
    return g
print(fancy_func(1, 2, 3, 4))
10

里面出现了神奇的两个函数 compile()exec()

  • compile(source, filename, mode[, flags[, dont_inherit]])
    • source:字符串或者 AST(Abstract Syntax Trees)对象
    • filename:代码文件名称,如果不是从文件读取代码则传递一些可辨认的值
    • mode:指定编译代码的种类。可以指定为 exec, eval, single
    • flags:变量作用域,局部命名空间,如果被提供,可以是任何映射对象
    • flags 和 dont_inherit 是用来控制编译源码时的标志
  • exec(obj)
    • obj:要执行的表达式。

命令式(解释型)编程和符号式编程的区别如下:

  • 命令式编程更容易使用。在 Python 中,命令式编程的大部分代码都是简单易懂的。命令式编程也更容易调试,这是因为无论是获取和打印所有的中间变量值,或者使用 Python 的内置调试工具都更加简单;
  • 符号式编程运行效率更高,更易于移植。符号式编程更容易在编译期间优化代码,同时还能够将程序移植到与 Python 无关的格式中,从而允许程序在非 Python 环境中运行,避免了任何潜在的与 Python 解释器相关的性能问题。

12.1.2 混合式编程

PyTorch 是基于命令式编程并且使用动态计算图。为了能够利用符号式编程的可移植性和效率,开发人员思考能否将这两种编程模型的优点结合起来,于是就产生了 TorchScript。TorchScript 允许用户使用纯命令式编程进行开发和调试,同时能够将大多数程序转换为符号式程序,以便在需要产品级计算性能和部署时使用。

接下来假设已经定义好了一个网络比如 net=MLP(),那么可以使用 net = torch.jit.script(net) 代码使用 TorchScript:

  • torch.jit.script(obj, optimize=None, _frames_up=0, _rcb=None, example_inputs=None)
    • 编写一个函数或 nn.Module 脚本将检查源代码,使用 TorchScript 编译器将其编译为 TorchScript 代码,并返回 ScriptModuleScriptFunction。 TorchScript 本身是 Python 语言的一个子集,因此并非 Python 中的所有功能都有效,但我们提供了足够的功能来计算张量并执行依赖于控制的操作。有关完整指南,请参阅 TorchScript 语言参考
    • 编写字典或列表的脚本会将其中的数据复制到 TorchScript 实例中,随后可以通过引用在 Python 和 TorchScript 之间以零复制开销传递。
    • torch.jit.script() 可以为模块、函数、字典和列表用作函数,而且还可以被用作装饰器。
    • 返回:
      • 如果 objnn.Module,script 会返回一个 ScriptModule。返回的 ScriptModule 将与原来的 nn.Module 有相同的子模块和参数集合。
      • 如果 obj 是独立的函数,一个 ScriptFunction 将会返回。
      • 如果 obj 是字典,将会返回 torch._C.ScriptDict
      • 如果 obj 是列表,将会返回 torch._C.ScriptList

在使用上面转化成 TorchScript 的代码后,一个三层的多层感知机大约增快了 20%。而且,还可以方便地使用 net.save('filepath.pt') 来保存网络结构。众所周知,普通的 torch.save()/torch.load() 是不能在没有原本的模块类定义下读取模型的。但是在 TorchScript 中,接下来即使我们删除了原本的多层感知机的类以及衍生的实例,也可以通过 torch.jit.load('filepath.pt') 重新载入模型。当然也不排除是我没删干净

12.2 异步计算

PyTorch 使用了 Python 自己的调度器来实现不同的性能权衡。对 PyTorch 来说 GPU操 作在默认情况下是异步的。当调用一个使用 GPU 的函数时,操作会排队到特定的设备上,但不一定要等到以后才执行。这允许并行执行更多的计算,包括在 CPU 或其他 GPU 上的操作。

因此,了解异步编程是如何工作的,通过主动地减少计算需求和相互依赖,有助于我们开发更高效的程序。这能够减少内存开销并提高处理器利用率。下面测试一下 numpy(CPU) 和 PyTorch(GPU) 的速度。

# GPU计算热身
device = d2l.try_gpu()
a = torch.randn(size=(1000, 1000), device=device)
b = torch.mm(a, a)

with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = numpy.random.normal(size=(1000, 1000))
        b = numpy.dot(a, a)

with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
numpy: 1.0981 sec
torch: 0.0011 sec

默认情况下,GPU 操作在 PyTorch 中是异步的。强制 PyTorch 在返回之前完成所有计算,这种强制说明了之前发生的情况:计算是由后端执行,而前端将控制权返回给了 Python。

例如下面调用 torch.cuda.synchronize(device),这个函数等待在一个 CUDA 设备上所有核的所有流都完成。

with d2l.Benchmark():
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
    torch.cuda.synchronize(device)
Done: 0.0089 sec

广义上说,PyTorch 有一个用于与用户直接交互的前端(例如通过 Python),还有一个由系统用来执行计算的后端。用户可以用各种前端语言编写 PyTorch 程序,如 Python 和 C++。不管使用的前端编程语言是什么,PyTorch 程序的执行主要发生在 C++ 实现的后端。由前端语言发出的操作被传递到后端执行。后端管理自己的线程,这些线程不断收集和执行排队的任务。请注意,要使其工作,后端必须能够跟踪计算图中各个步骤之间的依赖关系。因此,不可能并行化相互依赖的操作。

当语句的结果需要被打印出来时,Python 前端线程将等待 C++ 后端线程完成结果计算。这种设计的一个好处是 Python 前端线程不需要执行实际的计算。因此,不管 Python 的性能如何,对程序的整体性能几乎没有影响。

练习题

(1)在CPU上,对本节中相同的矩阵乘法操作进行基准测试,仍然可以通过后端观察异步吗?

torch 观察不到异步现象,反倒是 numpy 可以观察到异步的现象。虽然 torch.cuda.synchronize(torch.device('cpu')) 会弹出报错,但是仍然可以使用以下两个代码来测试速度:

with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = numpy.random.normal(size=(1000, 1000))
        b = numpy.dot(a, a)
        
# time.sleep(5)
        
with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
numpy: 0.9737 sec
torch: 0.2859 sec
with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = numpy.random.normal(size=(1000, 1000))
        b = numpy.dot(a, a)
        
time.sleep(5)
        
with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
numpy: 0.9414 sec
torch: 0.2103 sec

经过多次尝试,可以发现 torch 的执行时间有明显差异,这说明有 numpy 有部分仍然占用设备的时候,已经开始对 torch 的矩阵乘法计时了。

而如果把这两个矩阵乘法的顺序反过来,numpy 的时间变化不大,因此 torch 几乎没有异步而 numpy 异步了。

最后,我发现 torch.cuda.synchronize() 直接调用不加参数就不会报错了。如果它的 device 参数为 None,那么它将使用 current_device 函数找出当前设备。

12.3 自动并行

深度学习框架会在后端自动构建计算图。利用计算图,系统可以了解所有依赖关系,并且可以选择性地并行执行多个不相互依赖的任务以提高速度。

通常情况下单个操作符将使用所有CPU或单个GPU上的所有计算资源。并行化对单设备计算机来说并不是很有用,而并行化对于多个设备就很重要了。

请注意,接下来的实验至少需要两个GPU来运行。

12.3.1 基于 GPU 的并行计算

测试一下两个 GPU 串行各执行 10 次矩阵乘法和并行各执行 10 次矩阵乘法的速度。

devices = d2l.try_all_gpus()
def run(x):
    return [x.mm(x) for _ in range(50)]

x_gpu1 = torch.rand(size=(4000, 4000), device=devices[0])
x_gpu2 = torch.rand(size=(4000, 4000), device=devices[1])
run(x_gpu1)
run(x_gpu2)  # 预热设备
torch.cuda.synchronize(devices[0])
torch.cuda.synchronize(devices[1])

with d2l.Benchmark('GPU1 time'):
    run(x_gpu1)
    torch.cuda.synchronize(devices[0])

with d2l.Benchmark('GPU2 time'):
    run(x_gpu2)
    torch.cuda.synchronize(devices[1])
GPU1 time: 1.5491 sec
GPU2 time: 1.4804 sec

删除两个任务之间的 torch.cuda.synchronize() 语句,系统就可以在两个设备上自动实现并行计算。

with d2l.Benchmark('GPU1 & GPU2'):
    run(x_gpu1)
    run(x_gpu2)
    torch.cuda.synchronize()
GPU1 & GPU2: 1.5745 sec

12.3.2 并行计算与通信

在许多情况下,我们需要在不同的设备之间移动数据,比如在CPU和GPU之间,或者在不同的GPU之间。

通过在 GPU 上计算,然后将结果复制回 CPU 来模拟这个过程。

def copy_to_cpu(x, non_blocking=False):
    return [y.to('cpu', non_blocking=non_blocking) for y in x]

with d2l.Benchmark('在GPU1上运行'):
    y = run(x_gpu1)
    torch.cuda.synchronize()

with d2l.Benchmark('复制到CPU'):
    y_cpu = copy_to_cpu(y)
    torch.cuda.synchronize()
在GPU1上运行: 1.6285 sec
复制到CPU: 2.5801 sec

在 GPU 仍在运行时就开始使用 PCI-Express 总线带宽来移动数据是有利的。在 PyTorch 中,to()copy_() 等函数都允许显式的 non_blocking 参数,这允许在不需要同步时调用方可以绕过同步。设置 non_blocking=True 以模拟这个场景。

with d2l.Benchmark('在GPU1上运行并复制到CPU'):
    y = run(x_gpu1)
    y_cpu = copy_to_cpu(y, True)
    torch.cuda.synchronize()
在GPU1上运行并复制到CPU: 1.9456 sec

12.5 多 GPU 训练

在多个 GPU 上并行总共分为三种:

  1. 网络并行
  2. 按层并行
  3. 数据并行

实际上,数据并行是最常用的方法。原书中也重点讨论了数据并行。

12.5.2 数据并行性

假设一台机器有 \(k\) 个 GPU。 给定需要训练的模型,虽然每个 GPU 上的参数值都是相同且同步的,但是每个 GPU 都将独立地维护一组完整的模型参数。

一般来说,\(k\) 个 GPU 并行训练过程如下:

  • 在任何一次训练迭代中,给定的随机的小批量样本都将被分成 \(k\) 个部分,并均匀地分配到 GPU 上;
  • 每个 GPU 根据分配给它的小批量子集,计算模型参数的损失和梯度;
  • \(k\) 个 GPU 中的局部梯度聚合,以获得当前小批量的随机梯度;
  • 聚合梯度被重新分发到每个 GPU 中;
  • 每个 GPU 使用这个小批量随机梯度,来更新它所维护的完整的模型参数集。

在实践中请注意,当在 \(k\) 个 GPU 上训练时,需要扩大小批量的大小为 \(k\) 的倍数,这样每个 GPU 都有相同的工作量,就像只在单个 GPU 上训练一样。 因此,在 16-GPU 服务器上可以显著地增加小批量数据量的大小,同时可能还需要相应地提高学习率。

12.5.4 数据同步

对于高效的多 GPU 训练,我们需要两个基本操作。首先,我们需要向多个设备分发参数并附加梯度(get_params)。如果没有参数,就不可能在 GPU 上评估网络。第二,需要跨多个设备对参数求和,也就是说,需要一个 allreduce 函数。

get_params() 函数定义如下:

def get_params(params, device):
    new_params = [p.to(device) for p in params]
    for p in new_params:
        p.requires_grad_()
    return new_params

假设现在有一个向量分布在多个 GPU 上,下面的 allreduce 函数将所有向量相加,并将结果广播给所有GPU。请注意,需要将数据复制到累积结果的设备,才能使函数正常工作。

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):
        data[i][:] = data[0].to(data[i].device)

12.5.5 数据分发

nn.parallel.scatter() 是一个简单的工具函数,将一个小批量数据均匀地分布在多个 GPU 上。用法如下所示:

data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)
input : tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]])
load into [device(type='cuda', index=0), device(type='cuda', index=1)]
output: (tensor([[0, 1, 2, 3, 4],
        [5, 6, 7, 8, 9]], device='cuda:0'), tensor([[10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]], device='cuda:1'))

12.5.6 训练

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    # 在每个GPU上分别计算损失
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:  # 反向传播在每个GPU上分别执行
        l.backward()
    # 将每个GPU的所有梯度相加,并将其广播到所有GPU
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce(
                [device_params[c][i].grad for c in range(len(devices))])
    # 在每个GPU上分别更新模型参数
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0]) # 在这里,我们使用全尺寸的小批量
     

与前几章中略有不同:训练函数需要分配 GPU 并将所有模型参数复制到所有设备。显然,每个小批量都是使用 train_batch 函数来处理多个 GPU。我们只在一个 GPU 上计算模型的精确度,而让其他 GPU 保持空闲,尽管这是相对低效的,但是使用方便且代码简洁。

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # 将模型参数复制到num_gpus个GPU
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_iter:
            # 为单个小批量执行多GPU训练
            train_batch(X, y, device_params, devices, lr)
            torch.cuda.synchronize()
        timer.stop()
        # 在GPU0上评估模型
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
    print(f'测试精度:{animator.Y[0][-1]:.2f},{timer.avg():.1f}秒/轮,'
          f'在{str(devices)}')

12.6 多 GPU 的简洁实现

12.6.1 DataParallel()

原书出现了一个有趣的函数 net = nn.DataParallel(net, device_ids=devices),这个函数可以说是本节的重点。

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0) 这个函数:

在模块的层级上实现了数据并行。

这个容器通过在批次维度中分块将输入拆分到指定设备,从而并行化给定模块的应用程序(其他对象将在每个设备上复制一次)。在前向传递中,模块在每个设备上被复制,每个副本处理一部分输入。在向后传递期间,来自每个副本的梯度被汇总到原始模块中。

批量大小应大于使用的 GPU 数量。

另外,PyTorch 推荐使用 nn.parallel.DistributedDataParallel() 来代替 nn.Parallel(),原因如下:

大多数涉及批量输入和多个 GPU 的用例应默认使用 DistributedDataParallel 来利用多个 GPU。

使用具有多处理功能的 CUDA 模型有一些重要的注意事项;除非注意准确地满足数据处理要求,否则您的程序很可能会出现不正确或未定义的行为。

建议使用 DistributedDataParallel,而不是 DataParallel 进行多 GPU 训练,即使只有一个设备。

DistributedDataParallelDataParallel 之间的区别是:DistributedDataParallel 使用多进程,其中为每个 GPU 创建一个进程,而 DataParallel 使用多线程。通过使用 multiprocessing,每个 GPU 都有自己的专用进程,这避免了 Python 解释器的 GIL 带来的性能开销。

如果您使用 DistributedDataParallel,您可以使用 torch.distributed.launch 实用程序来启动您的程序,请参阅第三方后端。

允许将任意位置和关键字输入传递到 DataParallel 中,但某些类型需要特殊处理。张量将依托指定的维度被分开(默认为 0)。元组、列表和字典类型将被浅拷贝。其他类型将在不同的线程之间共享,如果写入模型的正向传播,则可能会被破坏。

在运行此 DataParallel 模块之前,并行化模块必须在 device_ids[0] 上具有其参数和缓冲区。原因在于:在每个 forward 中,模块在每个设备上被复制,因此对 forward 中正在运行的模块的任何更新都将丢失。例如,如果模块有一个计数器属性,在每次转发时递增,它将始终保持初始值,因为更新是在转发后销毁的副本上完成的。但是,DataParallel 保证 device[0] 上的副本的参数和缓冲区将与基本并行化模块共享存储。因此,将记录对设备 [0] 上的参数或缓冲区的就地更新。例如,BatchNorm2dspectral_norm() 依赖于此行为来更新缓冲区。

当模块在 forward() 中返回一个标量时,此 wrapper 将返回一个长度等于数据并行中使用的设备数量的向量,其中包含每个设备的结果。

参数:

  • module (Module) – 要并行的模块
  • device_ids (list of python:int or torch.device) – CUDA 设备(默认:全部设备)
  • output_device (int or torch.device) – 输出的设备位置(默认:device_ids[0])

于是,原书中训练这段代码写成了这样:

def train(net, num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    def init_weights(m):
        if type(m) in [nn.Linear, nn.Conv2d]:
            nn.init.normal_(m.weight, std=0.01)
    net.apply(init_weights)
    # 在多个GPU上设置模型
    net = nn.DataParallel(net, device_ids=devices)
    trainer = torch.optim.SGD(net.parameters(), lr)
    loss = nn.CrossEntropyLoss()
    timer, num_epochs = d2l.Timer(), 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    for epoch in range(num_epochs):
        net.train()
        timer.start()
        for X, y in train_iter:
            trainer.zero_grad()
            X, y = X.to(devices[0]), y.to(devices[0])
            l = loss(net(X), y)
            l.backward()
            trainer.step()
        timer.stop()
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
    print(f'测试精度:{animator.Y[0][-1]:.2f},{timer.avg():.1f}秒/轮,'
          f'在{str(devices)}')

注意第 \(19\) 行中是把数据传到了 \(0\) 号 GPU 上,然后它就会自动切成 GPU 个数据块然后传过去了。

运行代码测试一下!首先是只使用 \(1\) 块 GPU 的代码:

train(net, num_gpus=1, batch_size=256, lr=0.1)
测试精度:0.90,222.1秒/轮,在[device(type='cuda', index=0)]

然后是使用 \(2\) 块 GPU 的代码:

train(net, num_gpus=2, batch_size=512, lr=0.2)
测试精度:0.87,111.8秒/轮,在[device(type='cuda', index=0), device(type='cuda', index=1)]

接近一倍的速度提升。原书中跑一轮居然只需要 \(10\) 秒左右,不禁令人感慨。

热门相关:无量真仙   战神   夫人你马甲又掉了   战神   学霸女神超给力