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

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

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

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

一、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.py
import torch
import torch.distributed as dist

gpus = [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.py
import torch
import argparse
import 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 = torch.device("cuda", args.local_rank)

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

torch.cuda.set_device(args.local_rank)

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

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