动手学深度学习——数据并行，nn.DataParallel

目录

手动实现

网络定义

数据同步

向GPU设备传输参数

跨多个设备对参数求和

数据分发

训练

简洁实现

网络

训练

假设一台机器有𝑘个GPU。 给定需要训练的模型，虽然每个GPU上的参数值都是相同且同步的，但是每个GPU都将独立地维护一组完整的模型参数。 例如，下图演示了在𝑘=2时基于数据并行方法训练模型。

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

• 在任何一次训练迭代中，给定的随机的小批量样本都将被分成𝑘个部分，并均匀地分配到GPU上。
• 每个GPU根据分配给它的小批量子集，执行自己的前向传播和反向传播，计算模型参数的损失和梯度。
• 将𝑘个GPU中的梯度聚合，以获得当前小批量的随机梯度。
• 聚合梯度被分发到每个GPU中。
• 每个GPU使用这个小批量随机梯度，更新它所维护的模型参数（初始模型参数一样，并且更新所用的梯度也一样，所以各个GPU的模型参数是一样的）。

在实践中请注意，当在𝑘个GPU上训练时，需要扩大小批量的大小为𝑘的倍数，这样每个GPU都有相同的工作量，就像只在单个GPU上训练一样。 因此，在16-GPU服务器上可以显著地增加小批量数据量的大小，同时可能还需要相应地提高学习率。 还请注意，批量规范化也需要调整，例如，为每个GPU保留单独的批量规范化参数。小批量数据量更大时，学习率也需要稍微提高一些。

%matplotlib inline
import torch
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l

手动实现

网络定义

# 初始化模型参数
scale = 0.01
W1 = torch.randn(size=(20, 1, 3, 3)) * scale
b1 = torch.zeros(20)
W2 = torch.randn(size=(50, 20, 5, 5)) * scale
b2 = torch.zeros(50)
W3 = torch.randn(size=(800, 128)) * scale
b3 = torch.zeros(128)
W4 = torch.randn(size=(128, 10)) * scale
b4 = torch.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# 定义模型
def lenet(X, params):
    h1_conv = F.conv2d(input=X, weight=params[0], bias=params[1])
    h1_activation = F.relu(h1_conv)
    h1 = F.avg_pool2d(input=h1_activation, kernel_size=(2, 2), stride=(2, 2))
    h2_conv = F.conv2d(input=h1, weight=params[2], bias=params[3])
    h2_activation = F.relu(h2_conv)
    h2 = F.avg_pool2d(input=h2_activation, kernel_size=(2, 2), stride=(2, 2))
    h2 = h2.reshape(h2.shape[0], -1)
    h3_linear = torch.mm(h2, params[4]) + params[5]
    h3 = F.relu(h3_linear)
    y_hat = torch.mm(h3, params[6]) + params[7]
    return y_hat

# 交叉熵损失函数
loss = nn.CrossEntropyLoss(reduction&#