多GPU训练
问题拆分
- 在多个GPU之间拆分⽹络
每个GPU将流⼊特定层的数据作为输⼊,跨多个后续层对数据进⾏处理,然后将数据发送到下⼀个GPU
除⾮存框架或操作系统本⾝⽀持将多个GPU连接在⼀起,否则不建议这种⽅法
- 拆分层内的⼯作
例如,将问题分散到4个GPU,每个GPU⽣成16个通道的数据,⽽不是在单个GPU上计算64个通道。对于全连接的层,同样可以拆分输出单元的数量。
因为每⼀层都依赖于所有其他层的结果。此外,需要传输的数据量也可能⽐跨GPU拆分层时还要⼤。因此,基于带宽的成本和复杂性,我们同样不推荐这种⽅法
- 跨多个GPU对数据进⾏拆分
所有GPU尽管有不同的观测结果,但是执⾏着相同类型的⼯作。在完成每个⼩批量数据的训练之后,梯度在GPU上聚合。这种⽅法最简单,并可以应⽤于任何情况,同步只需要在每个⼩批量数据处理之后进⾏。
总体⽽⾔,只要GPU的显存⾜够⼤,数据并⾏是最⽅便的。****
数据并行性
⼀般来说, k个GPU并⾏训练过程如下:
- 在任何⼀次训练迭代中,给定的随机的⼩批量样本都将被分成k个部分,并均匀地分配到GPU上;
- 每个GPU根据分配给它的⼩批量⼦集,计算模型参数的损失和梯度;
- 将k个GPU中的局部梯度聚合,以获得当前⼩批量的随机梯度; 聚合梯度被重新分发到每个GPU中;
- 每个GPU使⽤这个⼩批量随机梯度,来更新它所维护的完整的模型参数集。
数据同步
对于⾼效的多GPU训练,我们需要两个基本操作。
⾸先,我们需要向多个设备分发参数并附加梯度(get_params)
第⼆,需要跨多个设备对参数求和,也就是说,需要⼀个allreduce函数。
数据分发
我们需要⼀个简单的⼯具函数,将⼀个⼩批量数据均匀地分布在多个GPU上。
多GPU简洁实现
import torch
from torch import nn
from d2l import torch as d2l# 1. 简单残差18
def resnet18(num_classes, in_channels=1):"""稍加修改的ResNet-18模型"""def resnet_block(out_channels, num_residuals, first_block=False):blk = []for i in range(num_residuals):if i == 0 and not first_block:blk.append(d2l.Residual(out_channels, use_1x1conv=True, strides=2))else:blk.append(d2l.Residual(out_channels))return nn.Sequential(*blk)net = nn.Sequential(nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(64),nn.ReLU())net.add_module("resnet_b1", resnet_block(64, 2, first_block=True))net.add_module("resnet_b2", resnet_block(128, 2))net.add_module("resnet_b3", resnet_block(256, 2))net.add_module("resnet_b4", resnet_block(512, 2))net.add_module("global_avg_pool",nn.AdaptiveAvgPool2d((1, 1)))net.add_module("fc", nn.Sequential(nn.Flatten(),nn.Linear(512, num_classes)))return net# 2.网络初始化
net = resnet18(10)device = d2l.try_all_gpus()# 3.训练模型
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(), 10animator = 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)}')# 1 GPU训练
train(net, num_gpus=1, batch_size=256, lr=0.1)
# 2 GPU训练
# train(net, num_gpus=2, batch_size=512, lr=0.2)