pytorch单机多卡的正确打开方式 以及可能会遇到的问题和相应的解决方法（pytorch单机多卡训练）

pytorch单机多卡的正确打开方式 以及可能会遇到的问题和相应的解决方法（pytorch单机多卡训练）

pytorch 单机多卡的正确打开方式

pytorch 使用单机多卡，大体上有两种方式

简单方便的 torch.nn.DataParallel(很 low，但是真的很简单很友好)使用torch.distributed加速并行训练(推荐，但是不友好)

首先讲一下这两种方式分别的优缺点

nn.DataParallel优点：就是简单缺点就是：所有的数据要先load到主GPU上，然后再分发给每个GPU去train，注意这时候主GPU的显存占用很大，你想提升​​batch_size​​​,那你的主GPU就会限制你的​​batch_size​​​，所以其实多卡提升速度的效果很有限注意：模型是会被copy到每一张卡上的，而且对于每一个BATCH的数据，你设置的​​batch_size​​​会被分成几个部分，分发给每一张卡，意味着，​​batch_size​​​最好是卡的数量​​n​​​的倍数，比如​​batch_size=6​​​,而你有​​n=4​​​张卡，那你实际上代码跑起来只能用​​3​​张卡，因为6整除3torch.distributed优点：避免了nn.DataParallel的主要缺点，数据不会再分发到主卡上，所以所有卡的显存占用很均匀缺点：不友好，调代码需要点精力，有很多需要注意的问题，我后面会列出

接下来展示如何使用两种方法以及相关注意事项

一、torch.nn.DataParallel

主要的修改就是用​​nn.DataParallel​​处理一下你的​​model​​

​​model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])​​

这个很简单，就直接上个例子，根据这个例子去改你的代码就好，主要就是注意对​​model​​的修改

注意model要放在主GPU上：​​model.to(device)​​

# main.pyimport torchimport torch.distributed as distgpus = [0, 1, 2, 3]torch.cuda.set_device('cuda:{}'.format(gpus[0]))train_dataset = ...train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)model = ...model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0]) #注意model要放在主GPU上optimizer = optim.SGD(model.parameters())for epoch in range(100): for batch_idx, (data, target) in enumerate(train_loader): images = images.cuda(non_blocking=True) target = target.cuda(non_blocking=True) ... output = model(images) loss = criterion(output, target) ... optimizer.zero_grad() loss.backward() optimizer.step()

二、torch.distributed加速

与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，只需要编写一份代码，torch 就会自动将其分配给多个进程，分别在多个 GPU 上运行。

要想把大象装冰箱，总共分四步！！

（1）要使用​​torch.distributed​​​，你需要在你的​​main.py(也就是你的主py脚本)​​中的主函数中加入一个参数接口：​​--local_rank​​

parser = argparse.ArgumentParser()parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')args = parser.parse_args()print(args.local_rank)

（2）使用 init_process_group 设置GPU 之间通信使用的后端和端口：

dist.init_process_group(backend='nccl')

（3）使用 DistributedSampler 对数据集进行划分：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

（4）使用 DistributedDataParallel 包装模型

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

举个栗子，参照这个例子去设置你的代码结构

# main.pyimport torchimport argparseimport torch.distributed as dist#（1）要使用`torch.distributed`，你需要在你的`main.py(也就是你的主py脚本)`中的主函数中加入一个**参数接口：`--local_rank`**parser = argparse.ArgumentParser()parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')args = parser.parse_args()#（2）使用 init_process_group 设置GPU 之间通信使用的后端和端口：dist.init_process_group(backend='nccl')torch.cuda.set_device(args.local_rank)#（3）使用 DistributedSampler 对数据集进行划分：train_dataset = ...train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)#（4）使用 DistributedDataParallel 包装模型model = ...model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])optimizer = optim.SGD(model.parameters())for epoch in range(100): for batch_idx, (data, target) in enumerate(train_loader): images = images.cuda(non_blocking=True) target = target.cuda(non_blocking=True) ... output = model(images) loss = criterion(output, target) ... optimizer.zero_grad() loss.backward() optimizer.step()

然后，使用以下指令，执行你的主脚本，其中​​--nproc_per_node=4​​表示你的单个节点的GPU数量：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

问题来了！！

你可能会在完成代码之后遇到各种问题，我这里列举一些要注意的点，去避坑

如果你遇到的莫名奇妙报错的问题，尝试这样去修改你的代码

​​device​​ 的设置你需要设置一个​​​device​​​参数，用来给你的数据加载到GPU上，由于你的数据会在不同线程中被加载到不同的GPU上，你需要传给他们一个参数​​device​​​,用于​​a.to(device)​​​的操作（a是一个tensor）​​​device​​如下设置

device = torch.device("cuda", args.local_rank)

你也可以通过设置当前​​cuda​​​，使用​​a.cuda()​​把张量放到GPU上，但是不推荐，可能会有一些问题

torch.cuda.set_device(args.local_rank)

find_unused_parameters=True这个是为了解决你的模型中定义了一些在​​​forward​​​函数中没有用到的网络层，会被视为“unused_layer”,这会引发错误，所以你在使用 DistributedDataParallel 包装模型的时候，传一个​​find_unused_parameters=True​​的参数来避免这个问题，如下：

encoder=nn.parallel.DistributedDataParallel(encoder, device_ids=[args.local_rank],find_unused_parameters=True)

num_workers很好理解，尽量不要给你的​​​DataLoader​​​设置​​numworkers​​参数，可能会有一些问题（不要太强迫症）shuffle=False你的​​​DataLoader​​​不要设置​​shuffle=True​​

valid_loader = torch.utils.data.DataLoader( part_valid_set, batch_size=BATCH, shuffle=False, num_workers=num_workers,sampler=valid_sampler)

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。