多GPU训练
多GPU训练的两种方式
将数据集合拆分到不同的GPU上面进行训练
实现代码:
- 好处:
- 切割比较方便
- 比较容易实现,不会出现错误
- 训练速度加快
数据集合的切割主要依赖于两个函数
- allReduce 函数
- spilte 函数
split函数的作用是将不同的数据和标签进行拆分,并且将拆分完的数据分发到不同的GPU上面,实现分发的功能
allReduce函数,因为数据虽然在不同的GPU上面进行训练但是,所有的模型公用一组参数
每次在不同的GPU上训练一次需要将参数的调整Reduce自动同步到不同的GPU上面
深度学习耗时的关键在于计算模型的梯度,多GPU训练可以将计算梯度的任务分发到不同的GPU上面
决定深度学习的在于GPU的内存,一个模型如果层数够多那么他的参数是非常庞大的,保存参数以及参数梯度需要巨大的显存,真正影响GPU计算性能的不是GPU的计算速度,而是他的显存,从内存上读取数据的时间远远大于计算的时间
- allReduce 函数
def allreduce(data): # allreduce 函数就是将所有向量相加在广播到所有的gpu上
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) # 再将参数传播到各个gpu上
- split 函数
def split_batch(X, y, devices):
"""将X和y拆分到多个设备上"""
assert X.shape[0] == y.shape[0]
return (nn.parallel.scatter(X, devices), # 将标签和数据同时拆分到不同的gpu上
nn.parallel.scatter(y, devices)) # nn.parallel函数是一个从头到尾的平分函数而且数据每次平分都是按照统一规则 避免了 训练数据x 和标签不在同一gpu上的可能
计算函数的主要步骤
1.将数据和标签切割成不同的大小分发到不同的gpu上面
2.在每个GPU上分别计算损失
3.在所有gpu上面计算反向传播计算参数梯度
5.在每个GPU上分别更新模型参数 可以通多sgd或者其他梯度下降的方法来计算模型参数
4.将每个GPU的所有梯度相加,并将其广播到所有GPU
这里可以看到在面对计算数据的时候与以往计算和更新梯度的方式有所不同,因为这里面对一次参数模型对所有的数据和标签进行计算并将结果累加,然后计算梯度,并进行梯度下降
在原来的梯度下降算法中每一对标签和数据都会对模型的参数进行影响
计算函数的代码:
# 每个batch的训练函数
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]) # 在这里,我们使用全尺寸的小批量
# 训练函数
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)}')
多GPU的计算格式
多GPU训练模型中单次训练的最小的量是一个batch 斤斤计较ize的训练数据,训练完一个batch size统一对模型参数进行更新