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整除3
• torch.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.py import torch import 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.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 的設置
  你需要設置一個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) 創作挑戰賽新人創作獎勵來咯，堅持創作打卡瓜分現金大獎

總結

以上是生活随笔為你收集整理的pytorch单机多卡的正确打开方式 以及可能会遇到的问题和相应的解决方法的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。