Pytorch多GPU訓練

阿新 • • 發佈：2021-01-23

Pytorch多GPU訓練

1. torch.nn.DataParallel

torch.nn.DataParallel()這個主要適用於單機多卡。個人一般比較喜歡在程式開始前，import包之後使用os.environ['CUDA_VISIBLE_DEVICES']來優先設定好GPU。例如要使用物理上第0,3號GPU只要在程式中設定如下：

os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

**注意：**如上限定物理GPU後，程式實際上的編號預設為device_ids[0]，device_ids[1]。就是說程式所使用的顯示卡編號實際上是經過了一次對映之後才會對映到真正的顯示卡編號上面的。所以device_ids

這個引數後續就不用再另行設定了。

batch_size設定

batch——size的大小應該大於所使用的GPU的數量。還應當是GPU個數的整數倍，這樣劃分出來的每一塊都會有相同的樣本數量。現batch_size = 原batch_size * num_GPUs

載入模型

model = nn.DataParallel(model)
model = model.cuda()

當然直接指定device_ids也可以：

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
model = model.cuda()

載入資料

inputs = 
 inputs.cuda()
labels = labels.cuda()

注意點

注意：nn.DataParallel(model)這句返回的已經不是原始的m了，而是一個DataParallel，原始的m儲存在DataParallel的module變數裡面。解決方法：

儲存的時候就取出原始model:

torch.save(model.module.state_dict(), path)

或者載入的時候用一個DataParallel載入，再取出原始模型：

model = nn.DataParallel(Resnet18())
model.load_state_dict(torch.load( 
path))
model = model.module

優化器

在訓練過程中，你的優化器同樣可以使用nn.DataParallel，如下兩行程式碼：

optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

# 優化器使用：
optimizer.step() --> optimizer.module.step()

Warning

UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars;
will instead unsqueeze and return a vector.

關於此的討論：

https://github.com/pytorch/pytorch/issues/9811

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

引數說明：

module – 要被並行的module
device_ids – CUDA裝置，預設為所有裝置。所以為了省事一般在前面加os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'
output_device – 輸出裝置（預設為device_ids[0]) 。所以所使用的0號卡，視訊記憶體佔用總是比較高。 負載不均衡很嚴重的話，建議使用DistributedDataParallel

此容器通過將mini-batch劃分到不同的裝置上來實現給定module的並行。在forward過程中，module會在每個裝置上都複製一遍，每個副本都會處理部分輸入。在backward過程中，副本上的梯度會累加到原始module上。

dataparallel只是資料input被分到不同卡上，模型還是隻在device0上的.首先各個卡只計算到loss，然後0號卡做loss平均，最後分發到各個卡上求梯度並進行引數更新。

Reference:

OPTIONAL: DATA PARALLELISM

PyTorch官方中文

pytorch 多 gpu 並行訓練

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

Pytorch多GPU訓練