Pytorch多GPU训练记录

单节点多GPU并行训练

DistributedDataParallel使用

简介：在单节点多卡的情况下DistributedDataParallel 比 torch.nn.DataParallel快。DistributedDataParallel实际上是每块GPU创建一个进程，并把数据平均分配到多块GPU上训练

步骤：

1. 设置CUDA_VISIBLE_DEVICES：

可以在运行命令时在命令行里设置：CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py

也可以在train.py代码中最开始设置：os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

2. 分布式初始化：

train.py中添加

 torch.distributed.init_process_group(backend="nccl")

3. 配置每个进程的GPU：

local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

如果代码中使用了argparse, 一定要在参数里面加上--local_rank：

parser.add_argument('--local_rank', help='experiment configuration filename', default="lib/config/360CC_config.yaml", type=int)

否则会报local_rank的错误。

4. 设置分布式model

model = crnn.get_crnn(config)

model.to(device)

model = torch.nn.parallel.DistributedDataParallel(module=model, device_ids=[local_rank], output_device=local_rank)

5. 设置分布式data：

    train_dataset = get_dataset(config)(config, is_train=True)
    sampler_train = DistributedSampler(train_dataset)
    train_loader = DataLoader(
        dataset=train_dataset,
        batch_size=config.TRAIN.BATCH_SIZE_PER_GPU,
        # shuffle=config.TRAIN.SHUFFLE,
        sampler=sampler_train,
        num_workers=config.WORKERS,
        pin_memory=config.PIN_MEMORY,
    )

 

注意，使用了sampler不能再使用shuffle参数。

6. 训练：

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py --cfg ./lib/config/CF_config.yaml

 

注意torch.distributed.launch为必须，用于分配local_rank

 

坑

如果之前都是单卡训练，现在使用DistributedDataParallel进行单机多GPU训练时，load之前的weights会报load_state_dict参数名不匹配的问题，原因是单GPU训练时，参数名中不会有module.的前缀，而单机多GPU训练时的每个参数前面都有module.前缀。所以解决方法是，对需要load的使用单卡训练的weights每一个参数key都加一个module.的前缀即可：

model.load_state_dict({'module.'+k:v for k,v in checkpoint['state_dict'].items()})

 

 

参考：

非常棒的帖子：https://zhuanlan.zhihu.com/p/113694038

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

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