多GPU训练实现
简介实现-利用pytorch框架
使用的模型是ResNet18
import torch
from torch import nn
from d2l import torch as d2l
def resnet18(num_classes, in_channels=1):
def resnet_block(in_channels, 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(in_channels, out_channels,
use_1x1conv=True, strides=2))
else:
blk.append(d2l.Residual(out_channels, 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_block1", resnet_block(64, 64, 2, first_block=True))
net.add_module("resnet_block2", resnet_block(64, 128, 2))
net.add_module("resnet_block3", resnet_block(128, 256, 2))
net.add_module("resnet_block4", resnet_block(256, 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
net = resnet18(10)
devices = d2l.try_all_gpus()
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, mean=0, 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
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()
print('$$$:', timer.avg(), str(devices))
train(net, num_gpus=1, batch_size=256, lr=0.1)
train(net, num_gpus=2, batch_size=256*2, lr=0.1)
从零实现单机多卡并行
使用的模型是Lenet
import torch
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l
scale = 0.01
# 20：表示张量的第一个维度的大小，这里代表卷积层的输出通道数（也就是卷积核的数量）
# 1：表示张量的第二个维度的大小，这里代表输入的通道数。假设输入图像是单通道（如灰度图），所以这里是 1
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)
# LeNet标准输入的图片尺寸为28x28，从28算下来就是50*4*4=50*16=800
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))
# 第0维表示批量大小 即有多少个样本这一批次
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='none')
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
new_params = get_params(params, d2l.try_gpu(0))
print('b1 weights device:', new_params[1].device)
print('b1 grad', new_params[1].grad)
def allreduce(data):
# 从索引 1 开始循环，直到 len(data) - 1，即不包括 0。
for i in range(1, len(data)):
# 要将数据搬到同一gpu设备上才能计算
# [:] 表示 对 data[0] 进行原地修改； 任何对 data[0][:] 的修改都会直接修改 data[0] 的内容，而不会创建一个新的副本
data[0][:] += data[i].to(data[0].device)
for i in range(1, len(data)):
data[i][:] = data[0].to(data[i].device)
data = [torch.ones((1, 2), device=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('allreduce之前： \n', data[0], '\n', data[1])
allreduce(data)
print('allreduce之后： \n', data[0], '\n', data[1])
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)
def split_batch(X, y, devices):
assert X.shape[0] == y.shape[0]
return (nn.parallel.scatter(X, devices),
nn.parallel.scatter(y, devices))
def train_batch(X, y, device_params, devices, lr):
X_shareds, y_shareds = split_batch(X, y, devices)
ls = [loss(lenet(X_shared, device_W), y_shared).sum()
# 各个GPU都有同一份全部的相关参数 存放在device_params之中，
# 而每次只取其中一个GPU设备的X_shared, y_shared 以及全部参数device_W
# 其实这个device_W在参数更新之前 每次循环都是一样的
for X_shared, y_shared, device_W in zip(
X_shareds, y_shareds, device_params)]
for l in ls:
l.backward()
with torch.no_grad():
# device_params[0]就是取第0行，len(device_params[0])就是第0行的长度也就是列数 表示网络参数的层数
# 就是网络有多少层，每一层都有自己的参数
for i in range(len(device_params[0])):
allreduce(
# c表示第c个gpu设备，对第i层的所有gpu设备的参数执行allreduce操作
[device_params[c][i].grad for c in range(len(devices))]
)
# 表示每个gpu设备上的参数
for param in device_params:
# 使用的是X.shape[0] 即整个批量的大小
d2l.sgd(param, lr, X.shape[0])
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)]
devices_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:
train_batch(X, y, devices_params, devices, lr)
torch.cuda.synchronize()
timer.stop()
# print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，'
# f'在{str(devices)}')
print('@@@:', timer.avg(), str(devices))
train(num_gpus=1, batch_size=256, lr=0.2)
train(num_gpus=2, batch_size=512, lr=0.3)
# 输出结果：
# @@@: 2.063412833213806 [device(type='cuda', index=0)]
# @@@: 1.7567662000656128 [device(type='cuda', index=0), device(type='cuda', index=1)]
