pytorch 分布式多卡训练DistributedDataParallel 踩坑记

ゝ一世哀愁。 2023-03-01 09:07 78阅读 0赞

### 前言： ###

近几天想研究了多卡训练，就花了点时间，本以为会很轻松，可是好多坑，一步一步踏过来，一般分布式训练分为单机多卡与多机多卡两种类型；

主要有两种方式实现：

１、DataParallel:   Parameter Server模式，一张卡位reducer，实现也超级简单，一行代码

DataParallel是基于Parameter server的算法，负载不均衡的问题比较严重，有时在模型较大的时候（比如bert-large），reducer的那张卡会多出3-4g的显存占用，也就是说用到DataParallel的时候gpu 0的使用内存会大大超出其他显卡的使用内存。

２、DistributedDataParallel：官方建议用新的DDP，采用all-reduce算法，本来设计主要是为了多机多卡使用，但是单机上也能用

为什么要分布式训练？

1.  可以用多张卡，总体跑得更快
2.  可以得到更大的 BatchSize
3.  有些分布式会取得更好的效果

分布式原理性解释：

[https://www.cnblogs.com/yh-blog/p/12877922.html][https_www.cnblogs.com_yh-blog_p_12877922.html]

[https://zhuanlan.zhihu.com/p/113694038][https_zhuanlan.zhihu.com_p_113694038]

主要分为以下几个部分：

1.  单机多卡，DataParallel（最常用，最简单）
2.  单机多卡，DistributedDataParallel（较高级）、多机多卡，DistributedDataParallel（最高级）
3.  如何启动训练
4.  模型保存与读取
5.  注意事项

### 一、单机多卡（DataParallel） ###

from torch.nn import DataParallel
    
    device = torch.device("cuda")
    ＃或者device = torch.device("cuda:0" if True else "cpu")
    
    model = MyModel()
    model = model.to(device)
    model = DataParallel(model)
    ＃或者model = nn.DataParallel(model,device_ids=[0,1，2,3])

比较简单，只需要加一行代码就行， model = DataParallel(model)

### 二、多机多卡、单机多卡（DistributedDataParallel） ###

建议先把注意事项看完在修改代码，防止出现莫名的bug，修改训练代码如下：

其中opt.local\_rank要在代码前面解析这个参数，可以去后面看我写的注意事项；

from torch.utils.data.distributed import DistributedSampler
        import torch.distributed as dist
        import torch
    
        # Initialize Process Group
        dist_backend = 'nccl'
        print('args.local_rank: ', opt.local_rank)
        torch.cuda.set_device(opt.local_rank)
        dist.init_process_group(backend=dist_backend)
    
        model = yourModel()＃自己的模型
        device = torch.device(opt.local_rank)
        model.to(device)＃经后来测试，好像注释这一条代码也可以正常训练；
        if torch.cuda.device_count() > 1:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            # 5) 封装
            # model = torch.nn.parallel.DistributedDataParallel(model,
            #                                                   device_ids=[opt.local_rank],
            #                                                   output_device=opt.local_rank)
            model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[opt.local_rank])
    
        dataset = ListDataset(train_path, augment=True, multiscale=opt.multiscale_training, img_size=opt.img_size, normalized_labels=True)#自己的读取数据的代码
        world_size = torch.cuda.device_count()
        datasampler = DistributedSampler(dataset, num_replicas=dist.get_world_size(), rank=opt.local_rank)
    
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=opt.batch_size,
            shuffle=False,
            num_workers=opt.n_cpu,
            pin_memory=True,
            collate_fn=dataset.collate_fn,
            sampler=datasampler
        )＃在原始读取数据中加sampler参数就行
    
    
    .....
    
    训练过程中，数据转cuda
          imgs = imgs.to(device)
          targets = targets.to(device)

### 三、如何启动训练 ###

**　１、DataParallel方式**

正常训练即可，即

python3  train.py

**       ２、DistributedDataParallel方式**

需要通过torch.distributed.launch来启动，一般是单节点，

CUDA_VISIBLE_DEVICES=0,1 python3 -m torch.distributed.launch --nproc_per_node=2 train.py

其中CUDA\_VISIBLE\_DEVICES　设置用的显卡编号，--nproc\_pre\_node 每个节点的显卡数量，一般有几个显卡就用几个显卡

当然还可以多节点训练，这个没了解清楚，想了解的可以戳这个

[https://zhuanlan.zhihu.com/p/113694038][https_zhuanlan.zhihu.com_p_113694038]

多节点

python３ -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=0
    ＃两个节点，在０号节点

要是训练成功，就会打印出几个信息，有几个卡就打印几个信息，如下图所示:

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2NvbXdheV9MaQ_size_16_color_FFFFFF_t_70][]

### 四、模型保存与读取 ###

以下a、b是对应的，用a保存，就用a方法加载

**　　１、保存**

a、只保存参数

torch.save(model.module.state_dict(), path)

b、保存参数与网络

torch.save(model.module,path)

**　　２、加载**

a、多卡加载模型预训练；

model = Yourmodel()
    if opt.pretrained_weights:
            if opt.pretrained_weights.endswith(".pth"):
                model.load_state_dict(torch.load(opt.pretrained_weights))
            else:
                model.load_darknet_weights(opt.pretrained_weights)

单卡加载模型，需要加载模型时指定主卡读模型，而且这个'cuda:0',是看你训练的模型是０还是１（否则就会出错RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device\_count() is 1. Please use torch.load with map\_location to map your storages to an existing device），可以根据自己的更改：

model = Yourmodel()
    if opt.pretrained_weights:
            if opt.pretrained_weights.endswith(".pth"):
                model.load_state_dict(torch.load(opt.pretrained_weights，map_location="cuda:0"))
            else:
                model.load_darknet_weights(opt.pretrained_weights)

b、单卡加载模型；

同样也要指定读取模型的卡。

model = torch.load(opt.weights_path, map_location="cuda:0")

多卡加载单卡的预训练模：

predict_params = torch.load(opt.pretrained_weights)
        model_dict = model.state_dict()
        for i,item in enumerate(model_dict.items()):
            name = item[0]
    
            model_dict[item[0]] = predict_params.state_dict()[name]
            model.load_state_dict(model_dict)

模型加载其他参考：[https://www.zhihu.com/question/67726969][https_www.zhihu.com_question_67726969]

### 五、注意事项 ###

**　１、model后面添加module**

获取到网络模型后，使用并行方法，并将网络模型和参数移到GPU上。**注意，若需要修改网络模块或者获得模型的某个参数，一定要在model后面加上.module，否则会报错，比如：**

model.img_size　　要改成　　model.module.img_size

**　　　２、.cuda或者.to(device)等问题**

device是自己设置，如果.cuda出错，就要化成相应的device

model（如：model.to(device)）  
　　　　input（通常需要使用Variable包装，如：input = Variable(input).to(device)）  
　　　　target（通常需要使用Variable包装  
　　　　nn.CrossEntropyLoss()（如：criterion = nn.CrossEntropyLoss().to(device)）  
　　　**３、args.local\_rank的参数**

通过torch.distributed.launch来启动训练，torch.distributed.launch 会给模型分配一个args.local\_rank的参数，所以在训练代码中要解析这个参数，也可以通过torch.distributed.get\_rank()获取进程id。

；

parser.add_argument("--local_rank", type=int, default=-1, help="number of cpu threads to use during batch generation")

** 　　　４、报错：`RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1`**

*  原因：在使用Pytorch加载模型时报错。加载的模型是用两个GPU训练的，而加载模型的电脑只有一个GPU，所以会出错。
 *  解决：`model = torch.load(model_path)`  
    改为：`model = torch.load(model_path, map_location='cuda:0')`  
    如果是4块到2块：就把`map_location`改为：`map_location={'cuda:1': 'cuda:0'}`

**　　　５、报错　RuntimeError: Address already in use**

原因：问题在于，TCP的端口被占用，一种解决方法是，运行程序的同时指定端口，端口号随意给出：

解决办法：１、在训练脚本里加　　--master\_port 29501

２、另一种方式，查找占用的端口号（在程序里 插入print输出），然后找到该端口号对应的PID值：`netstat -nltp`，然后通过`kill -9 PID`来解除对该端口的占用

[https_www.cnblogs.com_yh-blog_p_12877922.html]: https://www.cnblogs.com/yh-blog/p/12877922.html
[https_zhuanlan.zhihu.com_p_113694038]: https://zhuanlan.zhihu.com/p/113694038
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2NvbXdheV9MaQ_size_16_color_FFFFFF_t_70]: /images/20230209/ddb41d27e4074b05848bd00dab051431.png
[https_www.zhihu.com_question_67726969]: https://www.zhihu.com/question/67726969