单机多GPU训练

1 多GPU训练详细版

1.1 流程

1.基本思想

已知单机多GPU，把模型参数放在每个GPU上，把一个batch数据均分到每个GPU前向传播得到预测值，汇总每个GPU的损失得到总损失，用总损失计算梯度更新参数。

2. 流程

1）加载数据和设备（GPU）。
2）把初始化的参数复制到每个GPU上。
3）迭代训练，把一个小批量数据再多GPU训练。

1.2 知识点

1.参数处理

把参数放在GPU上，并设置可以求导。

def get_params(params,device):
    new_params = [p.clone().to(device) for p in params] # 参数默认是在cpu上处理，这个函数是把参数放在GPU上处理
    for p in new_params:
        p.requries_grad_()  # 计算参数的梯度
    return new_params

if __name__=="__main__":
    new_params = get_params(params,d2l.try_gpu(0))
    print("b1 weight:",new_params[1])
    print("b1 grad:",new_params[1].grad)

2.数据处理

把数据（特征、标签）拆分到多个GPU上。

def split_batch(X,y,devices):
    "将X,y拆分到多个设备上"
    assert X.shape[0] == y.shape[0]  # 判断特征和标签的样本数是否一致
    return (nn.parallel.scatter(X,devices),nn.parallel.scatter(y,devices))

if __name__=="__main__":
    data = torch.arange(20).reshape(4,5)
    devices = [torch.device("cuda:0"),torch.device("cuda:1")]
    split = nn.parallel.scatter(data,devices)
    

3.损失处理

把单个GPU的损失汇总，再把损失的汇总分给每个GPU用于计算题度。

##allreduce函数将所有向量相加，并将结果广播给所有的GPU
def allreduce(data):
    for i in range(1,len(data)):
        data[0][:] += data[i].to(data[0].device)  # 把data[i]的数据都给data[0]，得到总的数据
    for i in range(1,len(data)):
        data[i]+= data[0].to(data[i].device) # 把上一步总的数据给每个data[i]。

if __name__=="__main__":
    # 假设有两个GPU 
    data = [torch.ones((1,2),device=d2l.try_gpu(i))*(i+1) for i in range(2)]
    allreduce(data)
    #data = [[1,1]] , [[2,2]]-->
    #data = [[3,3]] , [[3,3]]

4.一个小批量上实现多GPU训练

1）把数据均分在每个GPU上。
2）计算每个GPU上的损失。
3）汇总每个GPU上的梯度。
4）用优化算法更新参数。

def train_batch(X,y,device_params,devices,lr):
    # 1 把数据均分在每个GPU上
    X_shards,y_shards = split_batch(X,y,devices)
    # 2 计算每个GPU上的损失
    ls = [
        loss(lenet(X_shards,device_w),y_shards).sum for X_shard,y_shard,device_w in zip(X_shards,y_shards ,device_params)]
    for l in ls:
        l.backward()
    with torch.no_grad():
        for i in range(len(device_params[0])): # i是模型的第i层
            # 3 汇总每个GPU上的梯度
            allreduce([device_params[c][i] for c in range(len(devices))])
        for param in device_params:
            # 4 用优化算法更新参数
            d2l.sgd(param,lr,X.shape[0])  # 对每个GPU更新参数

1.3 代码

# 1. 多GPU训练
#核心思想：把batch数据分给每个GPU计算损失，然后把损失反#向传播更新参数，以此类推，循环往复。

## 1.1 用lenet实践
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([20,1,3,3])*scale
b1 = torch.zeros(20)
w2 = torch.randn([50,20,5,5])*scale
b2 = torch.zeros(50)
w3 = torch.randn([80,123])*scale
b3 = torch.zeros(128)
w4 = torch.randn([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])
    h1_activation = F.relu(h1_conv)
    h1 = F.avg_pool2d(input=h1_activation,kernel=)
    h2_conv = F.conv2d(input=h1,weight=params[0])
    h2_activation = F.relu(h2_conv)
    h2 = F.avg_pool2d(input=h2_activation,kernel=)
    ......
 
##处理参数,把参数放在GPU上，并设置可以求导  
def get_params(params,device):
    new_params = [p.clone().to(device) for p in params] # 参数默认是在cpu上处理，这个函数是把参数放在GPU上处理
    for p in new_params:
        p.requries_grad_()  # 计算参数的梯度
    return new_params
    
new_params = get_params(params,2l.try_gpu(0))
print("b1 weight:",new_params[1])
print("b1 grad:",new_params[1].grad)

##allreduce函数将所有向量想家，并将结果广播给所有的GPU
def allreduce(data):
    for i in range(1,len(data)):
        data[0][:] += data[i].to(data[0].device)  # 把data[i]的数据都给data[0]，得到总的数据
    for i in range(1,len(data)):
        data[i]+= data[0].to(data[i].device) # 把上一步总的数据给每个data[i]。

#假设有两个GPU 
data = [torch.ones((1,2),device=d2l.try_gpu(i))*(i+1) for i in range(2)]
allreduce(data)
data = [[1,1]] , [[2,2]]
data = [[3,3]] , [[3,3]]

## 将一个小批量数据均匀地分布在多个GPU上
data = torch.arange(20).reshape(4,5)
devices = [torch.device("cuda:0"),torch.device("cuda:1")]
split = nn.parallel.scatter(data,devices)

##
def split_batch(X,y,devices):
    "将X,y拆分到多个设备上"
    assert X.shape[0] == y.shape[0]
    return (nn.parallel.scatter(X,devices),
            nn.parallel.scatter(y,devices))

##在一个小批量上实现多GPU训练
def train_batch(X,y,device_params,devices,lr):
    X_shards,y_shards = split_batch(X,y,devices)
    ls = [
        loss(lenet(X_shards,device_w),y_shards).sum for X_shard,y_shard,device_w in zip(X_shards,y_shards ,device_params)]
    for l in ls:
        l.backward()
    with torch.no_grad():
        for i in range(len(device_params[0])): # i是模型的第i层
            allreduce([device_params[c][i] for c in range(len(devices))])
        for param in device_params:
            d2l.sgd(param,lr,X.shape[0])  # 对每个GPU更新参数

# 定义训练函数
def train(net,num_gpus,batch_size,lr):
    # 1 加载数据和设备（GPU）
    train_iter,test_iter = d2l.load_data_fation_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # 2 把初始化的参数复制到每个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:
            # 3 在一个小批量上实现多GPU训练
            train_batch(X,y,device_params,devices,lr)
            torch.cuda.synchronize()  # 每个GPU做完后，同步一次
        timer.stop()
        animator.add(epoch+1,(d2l.evaluate_accuracy)......)
            
## 问题：精度下降，速度也没提升。可能原因是batch，lr、模型、torch平台没有并行优化。
     

2 多GPU训练精简版

2.1 代码

基本思想：用nn.DataParallel(net,device_ids=devices)实现单机多GPU训练

# 分布式训练
def train(net,num_gpus,batch_size,lr):
    # 1 加载数据和设备（GPU）
    train_iter,test_iter = d2l.load_data_fation_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.weights,std=0.01)
    
    net.apply(init_weights)
    net = nn.DataParallel(net,device_ids=devices)
    trainer = torch.optim.SGD(net.parameters(),lr)
    loss = nn.CrossEntropyLoss()
    timer,num_epochs = d2l.Timer(),10
    animator = 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()
        .......
         
    
        

参考李沐大神