pytorch 多GPU數據並行化
如果一個模型太大，一張顯卡上放不下，或者batch size太大，一張卡放不下，那麼就需要用多塊卡一起訓練，這時候涉及到 nn. DataParallel 的實用。
這個模塊的作用，本質上來說，就是：
看一份實驗代碼：
import torch
import torch.nn as nn
from torch.nn.parallel._functions import Broadcast, ReduceAddCoalesced
def backward(ctx, *grad_outputs):
print(grad_outputs)
return (None,) + ReduceAddCoalesced.apply(ctx.input_device, ctx.num_inputs, *grad_outputs)
Broadcast.backward = backward
x = torch.tensor([5.0, -1.0], dtype=torch.float).cuda().view(-1, 1)
model = nn.parallel.DataParallel(nn.Linear(in_features=1, out_features=1, bias=False).cuda(), device_ids=[0,1])
model.module.weight.data.zero_()
model.module.weight.data.add_(1.0)
y = model(x)
label = torch.zeros(2, 1, dtype=torch.float).cuda()
loss = torch.sum((y - label)**2)
loss.backward()
print(model.module.weight.grad)
對應的完整計算圖為：
這裡我發現，在batch dimension上，各個設備上的梯度最後是加起來的，體現在 ReduceAddCoalesced 的應用上。
不使用 DataParallel的話，會發現，原來梯度在batch dimension上本來就是求和的：
import torch
import torch.nn as nn
x = torch.tensor([5.0, -1.0], dtype=torch.float).cuda().view(-1, 1)
model = nn.Linear(in_features=1, out_features=1, bias=False).cuda()
model.weight.data.zero_()
model.weight.data.add_(1.0)
y = model(x)
label = torch.zeros(2, 1, dtype=torch.float).cuda()
loss = torch.sum((y - label)**2)
loss.backward()
print(model.weight.grad)
這段代碼與上面用DataParallel的結果是一樣的。
所以一般來說，batch size越大，weight上面的梯度就越大，對應的learning rate就應該變現。即 batch size 和 learning rate 成反比。
可是，有一點疑問：梯度下降的定義裡面，由於求的是梯度的近似，所以都有一個取平均的操作：
現在的框架基本上都不求平均了？
pytorch裡面是，tensorflow裡面也是：代碼
推薦閱讀：
TAG:PyTorch | 深度學習（DeepLearning） | 編程 |