import torch
import torch.nn as nn
import ipdb
class DataParallelModel(nn.Module):
def __init__(self):
super().__init__()
self.block1 = nn.Linear(10, 20)
def forward(self, x):
x = self.block1(x)
return x
def data_parallel(module, input, device_ids, output_device=None):
if not device_ids:
return module(input)
if output_device is None:
output_device = device_ids[0]
replicas = nn.parallel.replicate(module, device_ids)
print(f"replicas:{replicas}")
inputs = nn.parallel.scatter(input, device_ids)
print(f"inputs:{type(inputs)}")
for i in range(len(inputs)):
print(f"input {i}:{inputs[i].shape}")
replicas = replicas[:len(inputs)]
outputs = nn.parallel.parallel_apply(replicas, inputs)
print(f"outputs:{type(outputs)}")
for i in range(len(outputs)):
print(f"output {i}:{outputs[i].shape}")
result = nn.parallel.gather(outputs, output_device)
return result
model = DataParallelModel()
x = torch.rand(16,10)
result = data_parallel(model.cuda(),x.cuda(), [0,1])
print(f"result:{type(result)}")
最后輸出為
replicas:[DataParallelModel(
(block1): Linear(in_features=10, out_features=20, bias=True)
), DataParallelModel(
(block1): Linear(in_features=10, out_features=20, bias=True)
)]
inputs:<class 'tuple'>
input 0:torch.Size([8, 10])
input 1:torch.Size([8, 10])
outputs:<class 'list'>
output 0:torch.Size([8, 20])
output 1:torch.Size([8, 20])
result: torch.Size([16, 20])
可以看到整個流程如下：
- replicas: 將模型復制若干份，這里只有兩個GPU，所以復制兩份
- scatter: 將輸入數據若干等分，這里划分成了兩份，會返回一個tuple。因為batch size=16，所以剛好可以划分成8和8，那如果是15怎么辦呢？沒關系，它會自動划分成8和7，這個你自己可以做實驗感受一下。
- parallel_apply: 現在模型和數據都有了，所以當然就是並行化的計算咯，最后返回的是一個list，每個元素是對應GPU的計算結果。
- gather：每個GPU計算完了之后需要將結果發送到第一個GPU上進行匯總，可以看到最終的tensor大小是[16,20],這符合預期。