最近訓練模型時候想要使用使用多GPU運算來提高計算速度雪隧，參考一些博客以及自己的動手實驗搞懂了Pytorch的Multi-GPU原理。現(xiàn)在稍微整理一下员舵。
原理
通常情況下脑沿，多GPU運算分為單機多卡和多機多卡，兩者在pytorch上面的實現(xiàn)并不相同马僻，因為多機時庄拇，需要多個機器之間的通信協(xié)議等設置。
pytorch實現(xiàn)單機多卡十分容易韭邓，其基本原理就是：加入我們一次性讀入一個batch的數(shù)據(jù), 其大小為[16, 10, 5]措近，我們有四張卡可以使用。那么計算過程遵循以下步驟：
- 假設我們有4個GPU可以用女淑，pytorch先把模型同步放到4個GPU中瞭郑。
- 那么首先將數(shù)據(jù)分為4份，按照次序放置到四個GPU的模型中鸭你，每一份大小為[4, 10, 5]屈张；
- 每個GPU分別進行前項計算過程；
- 前向過程計算完后袱巨，pytorch再從四個GPU中收集計算后的結(jié)果假設[4, 10, 5]阁谆，然后再按照次序?qū)⑵淦唇悠饋韀16, 10, 5]，計算loss愉老。
整個過程其實就是 同步模型參數(shù)→分別前向計算→計算損失→梯度反傳
實現(xiàn)
定義模型
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
self.linear = nn.Linear(5, 5)
def forward(self, inputs):
output = self.linear(inputs)
return output
model = Model()
定義優(yōu)化器
optimizer = torch.optim.SGD(model.parameters(), lr = 0.01, momentum=0.9)
實現(xiàn)多GPU
# 假設就一個數(shù)據(jù)
data = torch.rand([16, 10, 5])
# 前向計算要求數(shù)據(jù)都放進GPU0里面
# device = torch.device('cuda:0')
# data = data.to(device)
data = data.cuda()
# 將網(wǎng)絡同步到多個GPU中
model_p = torch.nn.DataParalle(model.cuda(), device_ids=[0, 1, 2, 3])
logits = model_p(inputs)
# 接下來計算loss
loss = crit(logits, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
注：模型要求所有的數(shù)據(jù)和初始網(wǎng)絡被放置到GPU0场绿，實際上并不需要，只需要保證數(shù)據(jù)和初始網(wǎng)路都在你所選擇的多個gpu中的第一塊上就行嫉入。
