李沐动手学深度学习V2-多GPU数据并行使用Pytorch框架简洁实现

1. 多GPU数据并行简洁实现

一般来说， 𝑘 个GPU数据并行训练过程如下：

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

Pytorch框架通过nn.DataParallel()多GPU数据并行函数来实现将网络多GPU数据并行，得知使用了多少个GPU,同时将小批量数据平均分配到所有GPU上，然后计算loss和梯度，最后再把每个GPU上面的梯度聚合在一起，然后广播到所有GPU上面，进行网络参数权重更新，一行代码包含了上一节手动实现多GPU数据并行训练网络全部步骤，也即是上面所有步骤，如下面代码所示。

#获取所有的GPU
devices = [d2l.torch.try_gpu(i) for i in range(num_gpu)]
#网络初始化
resnet = resnet18(num_classes=10,in_channels=1)
# 在多个GPU上对模型进行数据并行
resnet = nn.DataParallel(module=resnet,device_ids=devices)

1. 使用Pytorch框架的高级API实现多GPU数据并行，使用了比LeNet更复杂的网络ResNet18，跟之前使用的ResNet模型不同的是：开始时使用了更小的卷积核、步长和填充，而且删除了最大汇聚层。

def train(net,num_gpu,lr,batch_size,epochs):
    train_iter,test_iter = d2l.torch.load_data_fashion_mnist(batch_size)
    #获取所有GPU
    devices = [d2l.torch.try_gpu(i) for i in range(num_gpu)]
    #初始化网络权重参数
    def init_weight(m):
        if type(m) in [nn.Linear,nn.Conv2d]:
            nn.init.normal_(m.weight,mean=0,std=0.01)
    net.apply(init_weight)
    # 在多个GPU对模型进行数据并行
    net = nn.DataParallel(module=net,device_ids=devices)
    net = net.to(devices[0])
    optim = torch.optim.SGD(net.parameters(),lr=lr)
    loss = nn.CrossEntropyLoss()
    timer = d2l.torch.Timer()
    animator = d2l.torch.Animator('epoch','test acc',xlim=[1,epochs])
    for epoch in range(epochs):
        net.train()#网络用于训练，必须有，因为模型中有batch_normalization层，batch_norm层在训练和测试两种状态下的权重参数不一样
        timer.start()
        for X,y in train_iter:
            optim.zero_grad()
            X = X.to(devices[0])#将输入数据和label都复制到第一个GPU上面
            y = y.to(devices[0])
            # 模型通过上面nn.DataParallel()多GPU数据并行函数得知使用了多少个GPU,同时将小批量数据平均分配到所有GPU上，然后计算loss和梯度，最后再把每个GPU上面的梯度聚合在一起，然后广播到所有GPU上面，进行网络参数权重更新
            y_hat = net(X)
            ls = loss(y_hat,y)
            ls.backward()
            optim.step()
        timer.stop()
        net.eval()#模型用于测试，必须有，因为模型中有batch_normalization层
        animator.add(epoch+1,(d2l.torch.evaluate_accuracy_gpu(net,test_iter),))
    print(f'测试精度：{animator.Y[0][-1]:.2f},{timer.avg():.2f}秒/轮,在{str(devices)}')

2. 模型训练和测试结果（学习率lr为0.1，轮数为10轮，批量大小为256个随机样本）

#网络初始化
resnet = resnet18(num_classes=10,in_channels=1)
lr,epochs,batch_size=0.1,10,256
num_gpu= 2  #本地只有一个GPU
train(resnet,num_gpu,lr,batch_size,epochs)
'''
补充：
#使用2个GPU进行训练，由于ResNet18比LeNet模型更复杂，因此多GPU数据并行效果更好，性能更高，每轮训练完所花的时间比使用一个GPU训练时间少了将近一半
train(net, num_gpus=2, batch_size=512, lr=0.2)#由于使用两个GPU数据并行训练，因此批量大小需要为一个GPU训练时两倍，学习率也要为一个GPU训练时的两倍
'''

在单个GPU上训练和测试结果如下图所示:

2. 多GPU数据并行简洁实现全部代码

import d2l.torch
import torch
from torch import nn
"""稍加修改的ResNet-18模型"""
def resnet18(num_classes,in_channels=1):
    def Residul_Block(in_channels,out_channels,num_residul,first_residul_block=False):
        blk = []
        for i in range(num_residul):
            if i==0 and not first_residul_block:
                blk.append(d2l.torch.Residual(in_channels,out_channels,use_1x1conv=True,strides=2))
            else:
                blk.append(d2l.torch.Residual(out_channels,out_channels,use_1x1conv=False,strides=1))
        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_block1',Residul_Block(in_channels=64,out_channels=64,num_residul=2,first_residul_block=True))
    net.add_module('resnet_block2',Residul_Block(in_channels=64,out_channels=128,num_residul=2,first_residul_block=False))
    net.add_module('resnet_block3',Residul_Block(in_channels=128,out_channels=256,num_residul=2,first_residul_block=False))
    net.add_module('resnet_block4',Residul_Block(in_channels=256,out_channels=512,num_residul=2,first_residul_block=False))
    net.add_module('adaptivepool2d',nn.AdaptiveAvgPool2d((1,1)))
    net.add_module('Flatten',nn.Flatten())
    net.add_module('Linear',nn.Linear(in_features=512,out_features=num_classes))
    return net
print(resnet18(10))

def train(net,num_gpu,lr,batch_size,epochs):
    train_iter,test_iter = d2l.torch.load_data_fashion_mnist(batch_size)
    #获取所有GPU
    devices = [d2l.torch.try_gpu(i) for i in range(num_gpu)]
    #初始化网络权重参数
    def init_weight(m):
        if type(m) in [nn.Linear,nn.Conv2d]:
            nn.init.normal_(m.weight,mean=0,std=0.01)
    net.apply(init_weight)
    # 在多个GPU对模型进行数据并行，将网络net使用多GPU数据并行进行训练
    net = nn.DataParallel(module=net,device_ids=devices)
    net = net.to(devices[0])
    optim = torch.optim.SGD(net.parameters(),lr=lr)
    loss = nn.CrossEntropyLoss()
    timer = d2l.torch.Timer()
    animator = d2l.torch.Animator('epoch','test acc',xlim=[1,epochs])
    for epoch in range(epochs):
        net.train()#网络用于训练，必须有，因为模型中有batch_normalization层，batch_norm层在训练和测试两种状态下的权重参数不一样
        timer.start()
        for X,y in train_iter:
            optim.zero_grad()
            X = X.to(devices[0])#将输入数据和label都复制到第一个GPU上面
            y = y.to(devices[0])
            # 模型通过上面nn.DataParallel()多GPU数据并行函数得知使用了多少个GPU,同时将小批量数据平均分配到所有GPU上，然后计算loss和梯度，最后再把每个GPU上面的梯度聚合在一起，然后广播到所有GPU上面，进行网络参数权重更新
            y_hat = net(X)
            ls = loss(y_hat,y)
            ls.backward()
            optim.step()
        timer.stop()
        net.eval()#模型用于测试，必须有，因为模型中有batch_normalization层
        animator.add(epoch+1,(d2l.torch.evaluate_accuracy_gpu(net,test_iter),))
    print(f'测试精度：{animator.Y[0][-1]:.2f},{timer.avg():.2f}秒/轮,在{str(devices)}')
#网络初始化
resnet = resnet18(num_classes=10,in_channels=1)
lr,epochs,batch_size=0.1,10,256
num_gpu= 2 
train(resnet,num_gpu,lr,batch_size,epochs)

3. 多机器（集群）多GPU分布式训练过程（如下图所示）

1. 在每台机器上读取一组（不同的）批量数据，在多个GPU之间分割数据并传输到GPU的显存中。基 于每个GPU上的批量数据分别计算预测和梯度。
2. 来自一台机器上的所有的本地GPU的梯度聚合在一个GPU上（或者在不同的GPU上聚合梯度的某些部分）。
3. 每台机器的梯度被发送到其本地CPU中。
4. 所有的CPU将梯度发送到中央参数服务器中，由该服务器聚合所有梯度。
5. 然后使用聚合后的梯度来更新参数，并将更新后的参数广播回各个CPU中。
6. 更新后的参数信息发送到本地一个（或多个）GPU中。
7. 所有GPU上的参数更新完成。
   

4. 上一节链接

李沐动手学深度学习V2-多GPU数据并行内容和手动实现代码