Pytorch distributed 多卡并行载入模型操作

一、pytorch distributed 多卡并行载入模型

这次来介绍下如何载入模型。

目前没有找到官方的distribute 载入模型的方式，所以采用如下方式。

大部分情况下，我们在测试时不需要多卡并行计算。

所以，我在测试时只使用单卡。

from collections import ordereddict
device = torch.device("cuda")
model = dgcnn(args).to(device)  #自己的模型
state_dict = torch.load(args.model_path)    #存放模型的位置

new_state_dict = ordereddict()
for k, v in state_dict.items():
    name = k[7:] # remove `module.`
    new_state_dict[name] = v
    # load params
model.load_state_dict (new_state_dict)

二、pytorch distributedparallel进行单机多卡训练

one_导入库：

import torch.distributed as dist
from torch.utils.data.distributed import distributedsampler

two_进程初始化：

parser = argparse.argumentparser()
parser.add_argument('--local_rank', type=int, default=-1)
# 添加必要参数
# local_rank：系统自动赋予的进程编号，可以利用该编号控制打印输出以及设置device

torch.distributed.init_process_group(backend="nccl", init_method='file://shared/sharedfile',
rank=local_rank, world_size=world_size)

# world_size：所创建的进程数，也就是所使用的gpu数量
# （初始化设置详见参考文档）

three_数据分发：

dataset = datasets.imagefolder(datapath)
data_sampler = distributedsampler(dataset, rank=local_rank, num_replicas=world_size)
# 使用distributedsampler来为各个进程分发数据，其中num_replicas与world_size保持一致，用于将数据集等分成不重叠的数个子集

dataloader = dataloader(dataset, batch_size=batch_size, num_workers=1,drop_last=true, pin_memory=true, sampler=data_sampler)
# 在dataloader中指定sampler时，其中的shuffle必须为false，而distributedsampler中的shuffle项默认为true，因此训练过程默认执行shuffle

four_网络模型：

torch.cuda.set_device(local_rank)
device = torch.device('cuda:'+f'{local_rank}')
# 设置每个进程对应的gpu设备

d = model()
d = torch.nn.syncbatchnorm.convert_sync_batchnorm(d).to(device)
# 由于在训练过程中各卡的前向后向传播均独立进行，因此无法进行统一的批归一化，如果想要将各卡的输出统一进行批归一化，需要将模型中的bn转换成syncbn
   
d = torch.nn.parallel.distributeddataparallel(
d, find_unused_parameters=true, device_ids=[local_rank], output_device=local_rank)
# 如果有forward的返回值如果不在计算loss的计算图里，那么需要find_unused_parameters=true，即返回值不进入backward去算grad，也不需要在不同进程之间进行通信。

five_迭代：

data_sampler.set_epoch(epoch)
# 每个epoch需要为sampler设置当前epoch

six_加载：

dist.barrier()
d.load_state_dict(torch.load('d.pth'), map_location=torch.device('cpu'))
dist.barrier()
# 加载模型前后用dist.barrier()来同步不同进程间的快慢

seven_启动：

cuda_visible_devices=1,3 python -m torch.distributed.launch --nproc_per_node=2 train.py --epochs 15000 --batchsize 10 --world_size 2
# 用-m torch.distributed.launch启动，nproc_per_node为所使用的卡数，batchsize设置为每张卡各自的批大小

以上为个人经验，希望能给大家一个参考，也希望大家多多支持。