batch size 训练时间_PyTorch多机多卡训练大规模数据

曾经终败给现在 2023-01-02 04:29 174阅读 0赞

目前由于训练数据量太大，我们想要把一大份数据拆分成多份送给不同的训练进程，通过这种并行处理数据的方式来加快模型的训练速度。PyTorch中实现分布式的模块有THD/C10D/torch.multiprocessing/torch.distributed.rpc/torch.distributed/DataParallel(DP)/DistributedDataParallel(DDP)等，本文将重点讲解一下DDP原理并给出具体实例操作。

## **DistributedDataParallel(DDP)模块** ##

DDP基于torch.distributed的模块，我们需要熟悉Backend和initialization methods:

1) initialization methods

因为DDP是真正的分布式，可以使用多台机器来组成一次并行运算的任务，所以就需要一种方法来给大家传递一个信息——如何联系到其它机器上的进程？目前DDP模块支持3种initialization methods：TCP initialization/Shared file-system initialization/Environment variable initialization

我们模型训练时采用的是TCP initialization(需提前准备两台带GPU的可相互进行TCP通信的服务器)

2) Backend

既然能够在不同的进程间进行通信，那必然是依赖于一些IPC的通信机制，这些通信机制一般是由PyTorch之外的三方实现的。在distributed模块中，一共有                4种不同的IPC通信backend：TCP/MPI/Gloo/NCCL

![2f7278b5ad5bd1289b6d9a018994926a.png][]

我们模型训练时是NCCL作为backend(由于是GPU训练，使用原厂的库NCCL会非常快，这是GPU上优化做得最好的后端)

**分布式模型代码**

主体训练代码如下：

def main(args):if not torch.cuda.is_available():raise Exception("need gpu to train network!")if args.seed is not None:
            random.seed(args.seed)
            np.random.seed(args.seed)
            torch.manual_seed(args.seed)
            torch.cuda.manual_seed(args.seed)
            torch.cuda.manual_seed_all(args.seed)
            cudnn.deterministic = Truecudnn.benchmark = Truecudnn.enabled = True# 多机多卡设置if args.gpu is not None:
            warnings.warn('You have chosen a specific GPU. This will completely '
                          'disable data parallelism.')if args.dist_url == "env://" and args.world_size == -1:
            args.world_size = int(os.environ["WORLD_SIZE"])
        args.distributed = args.world_size > 1 or args.multiprocessing_distributed
        ngpus_per_node = torch.cuda.device_count()if args.multiprocessing_distributed:args.world_size = ngpus_per_node * args.world_sizeprint("args.world_size", args.world_size)print("args.gpu", args.gpu)mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))else:# Simply call main_worker functionmain_worker(args.gpu, ngpus_per_node, args)def main_worker(gpu, ngpus_per_node, args):
        start_time = time.time()# 1 初始化init_methodargs.gpu = gpuif args.gpu is not None:print("Use GPU: {} for training".format(args.gpu))if args.distributed:if args.dist_url == "env://" and args.rank == -1:
                args.rank = int(os.environ["RANK"])if args.multiprocessing_distributed:# For multiprocessing distributed training, rank needs to be the
                # global rank among all the processesargs.rank = args.rank * ngpus_per_node + gpuprint("gpu", gpu)print("ngpus_per_node", ngpus_per_node)print("args.rank", args.rank)
            dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,world_size=args.world_size, rank=args.rank)# 2 加载训练模型
        # 2.1 确定模型并加载参数model = se_resnext101_32x4d()# 2.2 替换模型参数(1000->2)model.last_linear = torch.nn.Linear(model.last_linear.in_features, args.num_classes)# 2.3 是否微调模型 requires_grad=Truemodel_feature_layers = [model.layer0.parameters(), model.layer1.parameters(), model.layer2.parameters(), model.layer3.parameters(), model.layer4.parameters()]for parameters in model_feature_layers:for param in parameters:
                param.requires_grad = args.requires_grad

# 引入sync    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(args.gpu)

`# 2.5 多机多卡模型加载if not torch.cuda.is_available():print('using CPU, this will be slow')elif args.distributed:if args.gpu is not None: torch.cuda.set_device(args.gpu) model.cuda(args.gpu)args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])print("args.batch_size", args.batch_size)print("args.workers", args.workers)else: model.cuda()model = torch.nn.parallel.DistributedDataParallel(model)elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu)else:if args.arch.startswith('alexnet') or args.arch.startswith('vgg'): model.features = torch.nn.DataParallel(model.features) model.cuda()else: model = torch.nn.DataParallel(model).cuda()# 3 准备训练数据`

`# 3.1 读取lmdb格式数据,生成dataset  `

`traindir = os.path.join(Config.data_root, 'train.lmdb')`

`valdir = os.path.join(Config.data_root, 'val.lmdb') train_dataset = ImageFolderLMDB( traindir,transforms.Compose([ transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225]),])) val_dataset = ImageFolderLMDB( valdir,transforms.Compose([ transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225]),]))# 3.2 由dataset生成dataloader格式if args.distributed:print("args.distributed") train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)else: train_sampler = Nonetrain_loader = torch.utils.data.DataLoader( train_dataset,batch_size=args.batch_size,shuffle=(train_sampler is None),num_workers=args.workers,pin_memory=True,sampler=train_sampler) valid_loader = torch.utils.data.DataLoader( val_dataset,batch_size=args.batch_size,shuffle=False,num_workers=args.workers,pin_memory=True)# 4 定义损失函数和优化器, 学习率衰减 # 4.1 损失函数 # criterion = nn.CrossEntropyLoss().cuda()criterion = FocalLoss(alpha=0.25, gamma=2, num_class=args.num_classes, smooth=0.1)# 4.2 优化器optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay)# 4.3 设置学习率步长衰减scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=0.1)# 4.4 是否选择混合精度加快模型训练 #if args.apex:`

`# model, optimizer = amp.initialize(model, optimizer, opt_level='O1')`

`# 6 开始模型训练for epoch in range(start_epoch, args.epochs + 1):print("epoch", epoch)# 6.3.1 模型训练过程if args.distributed: train_sampler.set_epoch(epoch) acc1, losses = train(train_loader, model, criterion, optimizer, scheduler, epoch, args)# 6.3.2 模型验证过程`

`acc1, throughput = validate(valid_loader, model)`

`if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0):# 6.3.3 保存最后一次训练的模型torch.save( { 'epoch': epoch,'acc1': acc1,'loss': losses,'lr': scheduler.get_lr()[0],'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'scheduler_state_dict': scheduler.state_dict(),}, os.path.join(args.checkpoints, 'latest.pth'))# 6.3.4 每训练一个epoch即保存一次模型torch.save(model.module.state_dict(), os.path.join(args.checkpoints, "{}-epoch{}-acc{}.pth".format(args.network, epoch, acc1)))if __name__ == '__main__': args = parse_args()# logger = get_logger(__name__, args.log)`

`main(args)`

由于SyncBN能够完美支持多卡训练，而普通BN在多卡模式下实际上就是单卡模式，我们在训练代码中引入SyncBatchNorm：

# 引入syncmodel = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(args.gpu)

如何给不同进程分配数据集的不重叠、不交叉部分？DistributedSampler的实现方式是不同进程会使用一个相同的随机数种子，这样shuffle出来的东西就能确保一致。每次epoch开始前都要调用一下sampler的set\_epoch方法，这样才能让数据集随机shuffle起来。

# 3.2 由dataset生成dataloader格式if args.distributed:print("args.distributed")
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)else:
        train_sampler = Nonetrain_loader = torch.utils.data.DataLoader(
        train_dataset,batch_size=args.batch_size,shuffle=(train_sampler is None),num_workers=args.workers,pin_memory=True,sampler=train_sampler)

`os.makedirs(args.checkpoints)`

`# 6.3 开始模型训练 # logger.info('start training')for epoch in range(start_epoch, args.epochs + 1):print("epoch", epoch)# 6.3.1 模型训练过程if args.distributed: train_sampler.set_epoch(epoch)` **模型启动训练**

分别在两台服务器上启动训练代码，机器之间采用的是tcp通信。其中world\_size指分布式训练的节点数(指有多少台服务器参与训练)，rank指分布式训练的节点顺序，dist\_backend指采用nccl作为反向传播的backend。若world\_size参数错误，两台服务器的终端会卡在初始化进程组那里不动。

![c79f2626d1153718fb9990526f5aa41c.png][]

![4cf10d1a956363aaef2cac74b0eb890e.png][]

启动之后，训练报错：

![f6ba78fd53139b99e17841873be4669d.png][]

可能的原因是使用的是CUDA设备，因此使用NCCL backend；又因为没有IB连接，因此需要设置如下环境变量来禁用IB转而使用以太网。将NCCL\_IB\_DISABLE设置为1来禁止使用InfiniBand，转而使用 IP；如果网络接口不能被自动发现，则手工设置NCCL\_SOCKET\_IFNAME；如果还有问题，就设置如下的NCCL变量来打印更多的log：

export NCCL_SOCKET_IFNAME=eth0
    export NCCL_IB_DISABLE=1
    export NCCL_DEBUG=INFO
    export NCCL_DEBUG_SUBSYS=ALL

再次启动服务器训练命令，终端报错：

![ab510bc9aacdae1aba0af00958fa6cfa.png][]

可能的原因是训练使用的数据格式为lmdb形式，在dataloader阶段模型直接进行\_\_init\_\_初始化造成的问题。解决方案是延迟lmdb的初始化，可以将初始化操作移到\_\_getitem\_\_中进行。以下代码分别是在\_\_init\_\_初始化与\_\_getitem\_\_初始化，问题解决的前提是\_\_getitem\_\_初始化需要预先知道训练集和测试集的大小。

class ImageFolderLMDB(data.Dataset):def __init__(self, db_path, transform=None, target_transform=None):self.db_path = db_pathself.env = lmdb.open(db_path, subdir=osp.isdir(db_path),readonly=True, lock=False,readahead=False, meminit=False)with self.env.begin(write=False) as txn:self.length = loads_pyarrow(txn.get(b'__len__'))self.keys = loads_pyarrow(txn.get(b'__keys__'))self.transform = transformself.target_transform = target_transformdef __getitem__(self, index):
            env = self.env# 新增try:with env.begin(write=False) as txn:
                    byteflow = txn.get(self.keys[index])
                unpacked = loads_pyarrow(byteflow)# load imgimgbuf = unpacked[0]
                buf = six.BytesIO()
                buf.write(imgbuf)
                buf.seek(0)

`img = Image.open(buf).convert('RGB')`

class ImageFolderLMDB(data.Dataset):def __init__(self, db_path, transform=None, target_transform=None):self.db_path = db_pathself.env = None
            if 'train' in self.db_path:self.length = 447571elif 'val' in self.db_path:self.length = 55948else:raise NotImplementedErrorself.transform = transformself.target_transform = target_transformdef _init_db(self):self.env = lmdb.open(self.db_path, subdir=os.path.isdir(self.db_path),readonly=True, lock=False,readahead=False, meminit=False)with self.env.begin(write=False) as txn:# self.length = txn.stat()['entries'] - 1self.length = pa.deserialize(txn.get(b'__len__'))self.keys = pa.deserialize(txn.get(b'__keys__'))def __getitem__(self, index):if self.env is None:self._init_db()# 新增env = self.envtry:with env.begin(write=False) as txn:
                    byteflow = txn.get(self.keys[index])
                unpacked = loads_pyarrow(byteflow)# load imgimgbuf = unpacked[0]
                buf = six.BytesIO()
                buf.write(imgbuf)
                buf.seek(0)
                img = Image.open(buf).convert('RGB')

在分类上，DDP属于Data Parallel。简单来讲，就是通过提高batch size来增加并行度。DDP通过Ring-Reduce的数据交换方法提高了通讯效率，并通过启动多个进程的方式减轻Python GIL的限制，从而提高训练速度。一般来说，DDP都是显著地比DP快，能达到略低于卡数的加速比(例如，四卡下加速3倍)

再次启动服务器训练成功。可以看到训练时GPU-Util的使用率几乎都占到了100%，训练很高效。

![f4576b19c83ab00a36a8bbd22baf9c75.png][]

![9c634dd06baa318343d704d975c69eb4.png][]

**强调一下其它训练小trick：lmdb/NVIDIA-Apex/NVIDIA-DALI/SSD，超大数据量模型训练最重要的是有足够的SSD！！！**

## **附录：DDP的基本原理** ##

假如我们有N张显卡，

1.  (缓解GIL限制)在DDP模式下，会有N个进程被启动，每个进程在一张卡上加载一个模型，这些模型的参数在数值上是相同的。
2.  (Ring-Reduce加速)在模型训练时，各个进程通过一种叫Ring-Reduce的方法与其他进程通讯，交换各自的梯度，从而获得所有进程的梯度；
3.  (实际上就是Data Parallelism)各个进程用平均后的梯度更新自己的参数，因为各个进程的初始参数、更新梯度是一致的，所以更新后的参数也是完全相同的。

DP模式是很早就出现的、单机多卡的、参数服务器架构的多卡训练模式，在PyTorch，即是：

model = torch.nn.DataParallel(model)

在DP模式中，总共只有一个进程(受到GIL很强限制)。master节点相当于参数服务器，其会向其他卡广播其参数；在梯度反向传播后，各卡将梯度集中到master节点，master节点对搜集来的参数进行平均后更新参数，再将参数统一发送到其他卡上。这种参数更新方式，会导致master节点的计算任务、通讯量很重，从而导致网络阻塞，降低训练速度。

## **    Python GIL的存在使得，一个python进程只能利用一个CPU核心，不适合用于计算密集型的任务。使用多进程，才能有效率利用多核的计算资源。** ##

而DDP启动多进程训练，一定程度地突破了这个限制。

![04ac9fc4b8ff417d0d01e76c31e5ff96.png][]

Ring-Reduce是一种分布式程序的通讯方法。其原理参考文章：https://www.zhihu.com/question/57799212/answer/612786337

(1)因为提高通讯效率，Ring-Reduce比DP的parameter server快。其避免了master阶段的通讯阻塞现象，n个进程的耗时是o(n)。

(2)各进程独立计算梯度。

(3)每个进程将梯度依次传递给下一个进程，之后再把从上一个进程拿到的梯度传递给下一个进程。循环n次(进程数量)之后，所有进程就可以得到全部的梯度了。

(4)每个进程只跟自己上下游两个进程进行通讯，极大地缓解了参数服务器的通讯阻塞现象！

**参考文章**

https://zhuanlan.zhihu.com/p/178402798

https://zhuanlan.zhihu.com/p/187610959

https://zhuanlan.zhihu.com/p/250471767

https://www.zhihu.com/question/57799212/answer/612786337

[2f7278b5ad5bd1289b6d9a018994926a.png]: /images/20221120/9ed15c7bc01442068dcf13c469823e43.png
[c79f2626d1153718fb9990526f5aa41c.png]: /images/20221120/d0fd6cb345da4d2bb4bf5a0d0bb33e8f.png
[4cf10d1a956363aaef2cac74b0eb890e.png]: /images/20221120/bd8665e1a394404d891944dea228daeb.png
[f6ba78fd53139b99e17841873be4669d.png]: /images/20221120/bfee2b273e2645e98a0d3e9c9fc32c48.png
[ab510bc9aacdae1aba0af00958fa6cfa.png]: /images/20221120/1674e442c5ae4e0eba6189c3838359bf.png
[f4576b19c83ab00a36a8bbd22baf9c75.png]: /images/20221120/7c9d9287a89a40478b85640b3c5a58e8.png
[9c634dd06baa318343d704d975c69eb4.png]: /images/20221120/42fcdfcce2b24e2e85f748992933d572.png
[04ac9fc4b8ff417d0d01e76c31e5ff96.png]: /images/20221120/e8881079f93b41eda70608063f694336.png