pytorch 单机多卡训练distributedDataParallel
pytorch单机多卡：从DataParallel到DistributedDataParallel
最近想做的实验比较多，于是稍微学习了一下和pytorch相关的加速方式。本人之前一直在使用DataParallel做数据并行，在今天浅浅的学了下apex之后，发现apex和DataParrallel并不兼容，由此开始了DistributedDataParallel的研究。至于在单机上DistributedDataParallel本身已经较DataParallel更优秀之类的内容，网上已经有较多详细的描述，这里就不再赘述。
单GPU无并行
我们要做的只有两件事：
告诉进程，可用的GPU的有些
把模型，损失函数，数据都放到GPU上
对于第一点，一般我采取两种不同的方式：
命令行在python前指定 CUDA_VISIBLE_DEVICES=0,1
在py文件内部，通过os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0.1’ 指定
第二点的模型和损失函数直接通过使用cuda()转化：
model = model.cuda()
Loss = Loss.cuda()
唯一需要注意的点在于，需要通过model.cuda()把model传到GPU上之后，再把model.parameters()传给optimizer
数据在train和validate的过程里的每个batch里转化
for data,label in dataloader:
data = data.cuda()
label = label.cuda()
DataParallel
DataParallel的使用比较简单，只需要在原本单GPU的基础上加上一句即可。
model = torch.nn.DataParallel(model)
model = model.cuda()
torch会在每个batch下，把数据分配给各个GPU进行计算。在这里，0号和1号GPU各分到batch-size/2的数据。每个GPU对分得的数据求完导数之后，会将导数传给主进程，主进程再进行参数的更新，然后把更新后的参数传给各个GPU。所以DataParallel多GPU的计算结果和单GPU的计算结果基本没有差别(每次更新用的数据个数等于batch-size).
DistributedDataParallel
DistributedDataParallel使用更复杂一些,这里只说单机的情况,涉及到多节点的一些参数和函数,只能说在我本人的机器上已经确认能够正常使用.
DistributedDataParallel的过程和DataParallel稍有不同,各子进程之间只进行导数的传播.具体的,以上面的设置为例,GPU0对分配给自己的数据求导,然后把该批数据的导数传给GPU1;GPU1对分配给自己的数据求导,然后把该批数据的导数传给GPU0,因此每个GPU都有该批次数据完整的导数,然后每个GPU均进行一次梯度下降,因为参数和梯度以及优化器都是一致的,所以每个GPU独立更新一步之后的参数仍然保持一致.
这里初看进行了重复的运算(每个GPU的参数更新是相同的),但是在DataParallel时,虽然只有一个主GPU进行参数的更新,但是其他GPU在此时只是在等待主GPU返回新的参数,所以这部分重复的计算不会导致运算时间的增加.
另一方面,DataParallel时传输数据用时为:各GPU传导数给主GPU,主GPU传更新后的参数给各GPU.而DistributedDataParallel只有GPU彼此间传导数这一步.
此时需要做的事情有:
一系列我没仔细研究的初始化设置
告诉每个进程使用GPU的index
在每个进程里把model和data放到对应的GPU上
为每个GPU生成对应的数据
初始化以及命令行参数
这一部分我不太了解,有兴趣的可以学习其他资料,比如pytorch 分布式训练 distributed parallel 笔记
我这里只描述一下我用了可以work的code.
在main()的开始,先初始化分布式通信:
torch.distributed.init_process_group(backend='nccl', init_method='env://')通过命令行参数来获得进程index
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int,default=0)
args = parser.parse_args()
如果原来的程序以及在使用argparse,则只需插入这里的第二行即可.
其中的参数’–local_rank’也不需要手动给出,会在下面通过命令行里辅助工具torch.distributed.launch自动给出:
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS main.py [--arg1 --arg2 ...]
torch.distributed.launch将给每个进程分配一个编号’–local_rank’,下面我们将使用这个编号为各进程指定GPU.
NUM_GPUS为上面设置的可使用GPU的个数,这个例子里为2.
各进程GPU设置,model的设置
这里讲两种实现方式:
通过device=torch.device('cuda:{}'.format(args.local_rank)),来得到之前说的os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0.1’ 的第args.local_rank个GPU.在之后涉及到GPU时,使用tensor.to(device)即可.
通过torch.cuda.set_device(args.local_rank)来设置该进程使用的GPU序号,之后直接用tensor.cuda()即可
为每个GPU生成对应的数据
在DataParallel时,我们时通过一个完整的Dataloader来产生每个batch的数据,然后再自动把数据分配给各个GPU.使用DistributedDataParallel时,却是通过调用DistributedSampler来直接为各进程产生数据
train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False,
num_workers=2, pin_memory=True, sampler=train_sampler,)
这里需要设置shuffle=False,然后在每个epoch前,通过调用train_sampler.set_epoch(epoch)来达到shuffle的效果.注意事项
和DataParalle不同,DistributedDataParallel每个进程上的数据个数都等于batch-size,这导致每一次梯度下降实际用的数据量为batch-size*gpu-num.虽然一般情况下,batch越大计算越快,但是大batch不一定有最好的泛化,对于以及在单GPU或者DataParalle下以及对batch大小精调过的模型,设置batch为之前的GPU个数分之一才能复现之前结果
避免写入log文件发生冲突,只在某一个进程里用print,logger.可以通过args.local_rank==0来选择进程
计算loss或者accuracy时,要汇总各进程的信息,这里给出一个简单例子
def reduce_tensor(tensor):
# sum the tensor data across all machines
dist.all_reduce(rt, op=dist.reduce_op.SUM)
return rt
output = model(input)
loss = Loss(ouput, label)
accuracy = Accuracy(ouput, label)
temp_tensor = torch.tensor([0]) # create a temp tensor to load batchsize which is scalar
temp_tensor[0] = input.size(0)
log_loss = reduce_tensor(loss.clone().detach_())
log_acc = reduce_tensor(accuracy.clone().detach_()*batch_size)
input_size = reduce_tensor(temp_tensor.clone().detach_())
torch.cuda.synchronize() # wait every process finish above transmission
loss_total += log_loss.item()
if args.local_rank == 0:
print('loss_total=',loss_total)
print('accuracy=',log_acc/input_size)
因为刚写代码的时候对DistributedDataParallel的细节还不熟悉,所以按各GPU可能input.size()不一样在进行处理,通过了上面新建tensor来传输标量的笨办法来取得各GPU结果的权重.
原文链接：https://blog.csdn.net/weixin_39718268/article/details/105021631