多GPU训练能够加快模型的训练速度,而且在单卡上不能训练的模型可以使用多个小卡达到训练的目的。
多GPU训练可以分为单机多卡和多机多卡这两种,后面一种也就是分布式训练——训练方式比较麻烦,而且要关注的性能问题也有很多,据网上的资料有人建议能单机训练最好单机训练,不要使用多机训练。本文主要对单机多卡训练的实现展开说明。
这是最简单最直接的方法,代码中只需要一句代码就可以完成单卡多GPU训练了。其他的代码和单卡单GPU训练是一样的。
import torch
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
参数:
import torch
net = torch.nn.Linear(100,1)
print(net)
print('---------------------')
net = torch.nn.DataParallel(net, device_ids=[0,3])
print(net)
‘’‘输出’‘’
Linear(in_features=10, out_features=1, bias=True)
---------------------
DataParallel(
(module): Linear(in_features=10, out_features=1, bias=True)
)
nn.DataParallel是怎么做的:
概括来说就是:DataParallel会自动帮我们将数据切分 load 到相应 GPU,将模型复制到相应 GPU,进行正向传播计算梯度并汇总。
还有一句话,官网中是这样描述的:
The parallelized module must have its parameters and buffers on device_ids[0] before running this DataParallel module.
意思是:在运行此DataParallel模块之前,并行化模块必须在device_ids [0]上具有其参数和缓冲区。在执行DataParallel之前,会首先把其模型的参数放在device_ids[0]上,一看好像也没有什么毛病,其实有个小坑。
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"
当你添加这两行代码后,那么device_ids[0]默认的就是第2号卡,你的模型也会初始化在第2号卡上了,而不会占用第0号卡了。
当然你要保证上面这两行代码需要定义在下面这两行代码之前,一般放在train.py中import一些package之后:
device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)
优化器同样可以使用nn.DataParallel,如下两行代码:
optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)
1)多GPU计算减少了程序运行的时间?
很多时候发现在进行多GPU运算时,程序花费的时间反而更多了,这其实是因为你的batch_size太小了,因为torch.nn.DataParallel()这个函数是将每个batch的数据平均拆开分配到多个GPU上进行计算,计算完再返回来合并。这导致GPU之间的开关和通讯过程占了大部分的时间开销。
我们可以使用 watch -n 1 nvidia-smi
这个命令来查看每1s各个GPU的运行情况,如果发现每个GPU的占用率均低于50%,基本可以肯定你使用多GPU计算所花的时间要比单GPU计算花的时间更长了。
2)如何保存和加载多GPU网络?
import torch
net = torch.nn.Linear(10,1) # 先构造一个网络
net = torch.nn.DataParallel(net, device_ids=[0,3]) #包裹起来
torch.save(net.module.state_dict(), './networks/multiGPU.h5') #保存网络
# 加载网络
new_net = torch.nn.Linear(10,1)
new_net.load_state_dict(torch.load(“./networks/multiGPU.h5”))
因为DataParallel实际上是一个nn.Module,所以我们在保存时需要多调用了一个net.module,模型和优化器都需要使用net.module来得到实际的模型和优化器。
3)为什么第一块卡的显存会占用的更多一些?
最后一个参数output_device一般情况下是省略不写的,那么默认就是在device_ids[0],也就是第一块卡上,也就解释了为什么第一块卡的显存会占用的比其他卡要更多一些。
4)直接使用nn.DataParallel的时候,训练采用多卡训练,会出现一个warning?
UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars;
will instead unsqueeze and return a vector.
说明:
关于这个问题在pytorch官网的issues上有过讨论,下面简单摘出一些:
A batch of 3 would be calculated on a single GPU and results
would be [0.3, 0.2, 0.8] and model that returns the loss would return 0.43.
If cast to DataParallel, and calculated on 2 GPUs, [GPU1 - batch 0,1], [GPU2 - batch 2]
- return values would be [0.25, 0.8] (0.25 is average between 0.2 and 0.3)
- taking the average loss of [0.25, 0.8] is now 0.525!
Calculating on 3 GPUs, one gets [0.3, 0.2, 0.8] as results and average is back to 0.43!
这么求平均loss确实有不合理的地方。那么有什么好的解决办法呢,可以使用size_average=False,reduce=True作为参数。每个GPU上的损失将相加,但不除以GPU上的批大小。然后将所有平行损耗相加,除以整批的大小,那么不管几块GPU最终得到的平均loss都是一样的。
pytorch贡献者也实现了这个loss求平均的功能,即通过gather的方式来求loss平均:
https://github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be
如果它们在一个有2个GPU的系统上运行,DP将采用多GPU路径,调用gather并返回一个向量。如果运行时有1个GPU可见,DP将采用顺序路径,完全忽略gather,因为这是不必要的,并返回一个标量。
这种方法旨在缓解nn.DataParallel方法GPU使用效率低的问题。
1)首先第一步就是要进行init_process_group的初始化,声明GPU的NCCL通信方式。
import torch
torch.distributed.init_process_group(backend='nccl')
2)其次,由于是多线程的,因此数据加载和模型加载也要做对应的修改如下:
train_data = ReadDataSet('train.tsv', args, sentences_count = None)
train_sample = torch.utils.data.distributed.DistributedSampler(train_data)
train_loader = DataLoader(dataset=train_data, batch_size=args.batch_size, shuffle=(train_sample is None), sampler=train_sample)
model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], find_unused_parameters=True) #多进程多GPU并行
1)首先就是代码使用bash脚本启动的时候是不一样的,要像下面这么定义:
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main_gpus.py \
# 后面加一些你要传入的参数
在 pyton关键字之前把可用显卡号用它CUDA_VISIBLE_DEVICES=0,1来指定;同时python关键字之后-m torch.distributed.launch --nproc_per_node=2 来指定 分布式启动和采用的节点数,也就是有几个显卡也就用几个节点。
2)其次就是日志信息的打印,在代码中直接打印的话就会打印nproc_per_node次指定的输出信息,这个时候就需要指定进程号。
if (step+1)%200 == 0 and args.local_rank==0:
print('Train Epoch[{}/{}],step[{}/{}],tra_acc{:.6f} %,loss:{:.6f}'.format(epoch,epochs,step,len(train_iter),two_pro_train_acc*100,two_pro_loss))
这样就只会打印进程为0的对应各种信息。
3)再次就是loss和准确率的合并,这里有多个线程,肯定就需要对一个batch中多个线程对应的不同loss和准确率进行合并。实现方式如下:
def reduce_tensor(tensor: torch.Tensor):
rt = tensor.clone()
dist.all_reduce(rt, op=dist.ReduceOp.SUM)
rt /= dist.get_world_size() # 总进程数
return rt
把各自的loss或者accuracy做分布式操作的加法,然后在求平均值。
4)最后,关于batch_size和lr的设置,这里一般可以采用batch_size = n*batch_size_base的方式;而lr = (1,n)*lr_base的方式。
import torch
from torch import nn
from unet.unet_transfer import UNet16, UNetResNet
from pathlib import Path
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset, random_split
import torch.nn.functional as F
from torch.autograd import Variable
import shutil
from data_loader import ImgDataSet
import os
import argparse
import tqdm
import numpy as np
import scipy.ndimage as ndimage
import torch.distributed as dist
class AverageMeter(object):
“”“Computes and stores the average and current value”“”
def init(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def create_model(type =‘vgg16’):
if type ‘vgg16’:
print(‘create vgg16 model’)
model = UNet16(pretrained=True)
elif type ‘resnet101’:
encoder_depth = 101
num_classes = 1
print(‘create resnet101 model’)
model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True)
elif type == ‘resnet34’:
encoder_depth = 34
num_classes = 1
print(‘create resnet34 model’)
model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True)
else:
assert False
model.eval()
return model
def adjust_learning_rate(optimizer, epoch, lr):
“”“Sets the learning rate to the initial LR decayed by 10 every 30 epochs”“”
lr = lr * (0.1 ** (epoch // 30))
for param_group in optimizer.param_groups:
param_group[‘lr’] = lr
def find_latest_model_path(dir):
model_paths = []
epochs = []
for path in Path(dir).glob(‘*.pt’):
if ‘epoch’ not in path.stem:
continue
model_paths.append(path)
parts = path.stem.split(‘_’)
epoch = int(parts[-1])
epochs.append(epoch)
if len(epochs) > 0:
epochs = np.array(epochs)
max_idx = np.argmax(epochs)
return model_paths[max_idx]
else:
return None
def train(train_loader, model, criterion, optimizer, validation, args):
latest_model_path = find_latest_model_path(args.model_dir)
best_model_path = os.path.join(*[args.model_dir, 'model_best.pt'])
if latest_model_path is not None:
state = torch.load(latest_model_path)
epoch = state['epoch']
model.load_state_dict(state['model'])
epoch = epoch
#if latest model path does exist, best_model_path should exists as well
assert Path(best_model_path).exists() == True, f'best model path {best_model_path} does not exist'
#load the min loss so far
best_state = torch.load(latest_model_path)
min_val_los = best_state['valid_loss']
print(f'Restored model at epoch {epoch}. Min validation loss so far is : {min_val_los}')
epoch += 1
print(f'Started training model from epoch {epoch}')
else:
print('Started training model from epoch 0')
epoch = 0
min_val_los = 9999
valid_losses = []
for epoch in range(epoch, args.n_epoch + 1):
adjust_learning_rate(optimizer, epoch, args.lr)
tq = tqdm.tqdm(total=(len(train_loader) * args.batch_size))
tq.set_description(f'Epoch {epoch}')
losses = AverageMeter()
model.train()
for i, (input, target) in enumerate(train_loader):
two_pro_loss = 0
input_var = Variable(input).cuda(args.local_rank, non_blocking=True)
target_var = Variable(target).cuda(args.local_rank, non_blocking=True)
masks_pred = model(input_var)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = target_var.view(-1)
loss = criterion(masks_probs_flat, true_masks_flat)
two_pro_loss += reduce_tensor(loss).item() # 有多个进程,把进程0和1的loss加起来平均
losses.update(two_pro_loss)
tq.set_postfix(loss='{:.5f}'.format(losses.avg))
tq.update(args.batch_size)
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
valid_metrics = validation(model, valid_loader, criterion)
valid_loss = valid_metrics['valid_loss']
valid_losses.append(valid_loss)
print(f'\tvalid_loss = {valid_loss:.5f}')
tq.close()
#save the model of the current epoch
epoch_model_path = os.path.join(*[args.model_dir, f'model_epoch_{epoch}.pt'])
torch.save({
'model': model.state_dict(),
'epoch': epoch,
'valid_loss': valid_loss,
'train_loss': losses.avg
}, epoch_model_path)
if valid_loss < min_val_los:
min_val_los = valid_loss
torch.save({
'model': model.state_dict(),
'epoch': epoch,
'valid_loss': valid_loss,
'train_loss': losses.avg
}, best_model_path)
def reduce_tensor(tensor: torch.Tensor):
rt = tensor.clone()
dist.all_reduce(rt, op=dist.ReduceOp.SUM)
rt /= dist.get_world_size() # 总进程数
return rt
def validate(model, val_loader, criterion):
losses = AverageMeter()
model.eval()
with torch.no_grad():
for i, (input, target) in enumerate(val_loader):
two_pro_loss = 0
input_var = Variable(input).cuda(args.local_rank, non_blocking=True)
target_var = Variable(target).cuda(args.local_rank, non_blocking=True)
output = model(input_var)
loss = criterion(output, target_var)
two_pro_loss += reduce_tensor(loss)
losses.update(loss.item(), input_var.size(0))
return {'valid_loss': losses.avg}
def save_check_point(state, is_best, file_name = ‘checkpoint.pth.tar’):
torch.save(state, file_name)
if is_best:
shutil.copy(file_name, ‘model_best.pth.tar’)
def calc_crack_pixel_weight(mask_dir):
avg_w = 0.0
n_files = 0
for path in Path(mask_dir).glob(‘.’):
n_files += 1
m = ndimage.imread(path)
ncrack = np.sum((m > 0)[:])
w = float(ncrack)/(m.shape[0]*m.shape[1])
avg_w = avg_w + (1-w)
avg_w /= float(n_files)
return avg_w / (1.0 - avg_w)
if name == ‘main’:
parser = argparse.ArgumentParser(description=‘PyTorch ImageNet Training’)
parser.add_argument(‘–n_epoch’, default=10, type=int, metavar=‘N’, help=‘number of total epochs to run’)
parser.add_argument(‘–lr’, default=0.001, type=float, metavar=‘LR’, help=‘initial learning rate’)
parser.add_argument(‘–momentum’, default=0.9, type=float, metavar=‘M’, help=‘momentum’)
parser.add_argument(‘–print_freq’, default=20, type=int, metavar=‘N’, help=‘print frequency (default: 10)’)
parser.add_argument(‘–weight_decay’, default=1e-4, type=float, metavar=‘W’, help=‘weight decay (default: 1e-4)’)
parser.add_argument(‘–batch_size’, default=4, type=int, help=‘weight decay (default: 1e-4)’)
parser.add_argument(‘–num_workers’, default=4, type=int, help=‘num_workers’)
parser.add_argument(‘–data_dir’,type=str, default=‘dataset’, help=‘input dataset directory’)
parser.add_argument(‘–model_dir’, type=str, default=‘model’, help=‘output model directory’)
parser.add_argument(‘–model_type’, type=str, required=False, default=‘vgg16’, choices=[‘vgg16’, ‘resnet101’, ‘resnet34’])
parser.add_argument(‘–local_rank’, type=int, default=-1)
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)
# os.environ['LOCAL_RANK'] = -1
dist.init_process_group(backend='nccl')
os.makedirs(args.model_dir, exist_ok=True)
DIR_IMG = os.path.join(args.data_dir, 'images')
DIR_MASK = os.path.join(args.data_dir, 'masks')
img_names = [path.name for path in Path(DIR_IMG).glob('*.jpg')]
mask_names = [path.name for path in Path(DIR_MASK).glob('*.jpg')]
print(f'total images = {len(img_names)}')
model = create_model(args.model_type)
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion = nn.BCEWithLogitsLoss().cuda()
# ori
channel_means = [0.485, 0.456, 0.406]
channel_stds = [0.229, 0.224, 0.225]
# dam
# channel_means = [0.595, 0.608, 0.604]
# channel_stds = [0.047, 0.047, 0.047]
train_tfms = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(channel_means, channel_stds)])
val_tfms = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(channel_means, channel_stds)])
mask_tfms = transforms.Compose([transforms.ToTensor()])
dataset = ImgDataSet(img_dir=DIR_IMG, img_fnames=img_names, img_transform=train_tfms, mask_dir=DIR_MASK, mask_fnames=mask_names, mask_transform=mask_tfms)
train_size = int(0.85*len(dataset))
valid_size = len(dataset) - train_size
train_dataset, valid_dataset = random_split(dataset, [train_size, valid_size])
train_sample = torch.utils.data.distributed.DistributedSampler(train_dataset)
valid_sample = torch.utils.data.distributed.DistributedSampler(valid_dataset)
train_loader = DataLoader(train_dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, sampler=train_sample)
valid_loader = DataLoader(valid_dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, sampler=valid_sample)
model.cuda(args.local_rank)
train(train_loader, model, criterion, optimizer, validate, args)
!!!!!!!你可能会运行出错,记得一定要按照下面的去运行代码!!!!!!!!!
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main_gpus.py \
# 后面加一些你要传入的参数
我自己的跑成功了,但是还是很慢,不过由原来的3小时变成52min一个epoch。
这里直接贴上 这位大佬 的对比吧。
多GPU batch_size = 16 lr=[1,n]*lr_base(方法:nn.DataParallel(model))
其实在每个forward中batch_size 是8;按照经验来说lr应该是要扩大相应的倍数的
总时间是140s,最高准确率是61.4%,相比单卡速度提升了48.6%,耗费时间减少了32.7%。
如果要采用单机多卡训练模型,无疑是采用nn.parallel.DistributedDataParallel这种方式,速度最快;有限时间内,训练效果最好。
单机多卡训练模型加速方式有采用混合精度和梯度累加等。这里只有梯度累加能够起加速作用的训练是多卡训练才能享受到的,单卡并不能加速。简单的分析就是,多卡训练需要一个梯度同步的过程,就是GPU之间在每一个batch的计算上都会进行通信,这个时间就会导致训练处于等待状态。而梯度累加就是变相增大batch_size,减小batch数目,从而减少GPU之间的通信,起到加速作用。当然梯度累加的代码实现也比较简单,正常的训练代码:
for i, (inputs, labels) in enumerate(training_set):
loss = model(inputs, labels) # 计算loss
optimizer.zero_grad() # 清空梯度
loss.backward() # 反向计算梯度
optimizer.step() # 更新参数
使用梯度累加:
for i, (inputs, labels) in enumerate(training_set):
loss = model(inputs, labels) # 计算loss
loss = loss / accumulation_steps # Normalize our loss (if averaged)
loss.backward() # 反向计算梯度,累加到之前梯度上
if (i+1) % accumulation_steps == 0:
optimizer.step() # 更新参数
model.zero_grad() # 清空梯度
当然这里的效果取决于模型的大小,模型越大收益越大。
下面是采用robert_large模型、2张3090显卡做文本2分类的一个速度(1W训练集和1K验证集):
单机多卡
--accumulation_steps 2
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 686.4237
tra_acc73.325002 %,dev_acc76.400002 %,best_dev_acc76.400002 %
*******************************************************************************
--accumulation_steps 5
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 578.8834
tra_acc73.329997 %,dev_acc75.500000 %,best_dev_acc76.100006 %
***
--accumulation_steps 10
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 579.5692
tra_acc71.015000 %,dev_acc75.400402 %,best_dev_acc77.300002 %
***
--accumulation_steps 20
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 613.6300s
tra_acc64.775002 %,dev_acc78.199997 %,best_dev_acc78.199997 %
***
--accumulation_steps 20
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 580.7058
tra_acc64.754999 %,dev_acc77.400002 %,best_dev_acc77.400002 %
***
--accumulation_steps 50
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 621.0073
tra_acc53.034997 %,dev_acc71.900002 %,best_dev_acc71.900002 %
***
--accumulation_steps 80
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 568.5933
tra_acc43.325001 %,dev_acc67.199997 %,best_dev_acc67.199997 %
***
--accumulation_steps 80
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 575.0746
tra_acc44.005001 %,dev_acc67.500000 %,best_dev_acc67.500000 %
***
--accumulation_steps 0
roberta_large textclassification task train and dev 2 epochs time is 718.4363s
tra_acc74.285001 %,dev_acc73.199997 %,best_dev_acc73.199997 %
***
--accumulation_steps 0
roberta_large textclassification task train and dev 2 epochs time is 694.9744
tra_acc74.559999 %,dev_acc74.000000 %,best_dev_acc74.000000 %
单卡单GPU
***
trian and eval model time is 1023.3577s
tra_acc64.715000 %,dev_acc71.400000 %,best_dev_acc71.400000 %
***
trian and eval model time is 1034.7063
tra_acc72.760000 %,dev_acc74.300000 %,best_dev_acc74.300000 %
***
结论:
单卡3090耗时:1029s
双卡3090耗时:707s——提升:45.5%
双卡3090+梯度累加耗时: 580s——提升77.4%,21.9%
可以看到在当前数据集和模型的情况下,accumulation_step = 10 可以取得最好的效果,相对于单卡提速77.4%;双卡梯度累加相对于双卡不采用梯度累加提速21.9%,前提是模型的准确率并没有降低。这个trick就很好用了。
【1】https://blog.csdn.net/HUSTHY/article/details/108226256
【2】https://zhuanlan.zhihu.com/p/86441879
【3】https://zhuanlan.zhihu.com/p/145427849
【4】https://blog.csdn.net/qq_38410428/article/details/119392993
【5】https://blog.csdn.net/wangkaidehao/article/details/104411682
【6】https://zhuanlan.zhihu.com/p/102697821