Pytorch并行计算:nn.parallel.replicate, scatter, gather, parallel_apply
2019-09-17 15:31
337 查看
原文链接:http://www.cnblogs.com/marsggbo/p/11534141.html
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)}")
[/code]
最后输出为
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])
[/code]
可以看到整个流程如下:
- replicas: 将模型复制若干份,这里只有两个GPU,所以复制两份
- scatter: 将输入数据若干等分,这里划分成了两份,会返回一个tuple。因为batch size=16,所以刚好可以划分成8和8,那如果是15怎么办呢?没关系,它会自动划分成8和7,这个你自己可以做实验感受一下。
- parallel_apply: 现在模型和数据都有了,所以当然就是并行化的计算咯,最后返回的是一个list,每个元素是对应GPU的计算结果。
- gather:每个GPU计算完了之后需要将结果发送到第一个GPU上进行汇总,可以看到最终的tensor大小是[16,20],这符合预期。
MARSGGBO♥原创
2019-9-17
转载于:https://www.cnblogs.com/marsggbo/p/11534141.html
相关文章推荐
- C#并行计算类Parallel的性能测试
- 并行计算介绍(Introduction to Parallel Computing ) ——我主张阅读英文原著
- python并行计算 [joblib] Attempting to do parallel computing without protecting。。。
- Parallel ForEach For 多线程并行计算使用注意
- 文献:利用自驱动分子马达并行计算子集和问题 Parallel computation with molecular-motor-propelled agents...(PNAS)
- 基于微软并行计算的工作流解决方案 (3) AsyncParallelActivity
- 并行计算—parallel,for联合使用
- 分布式并行计算方案:parallel computing by kafka-storm 发布了
- R语言并行计算的原理和案例(foreach,parallel)
- linux命令parallel实现多进程并行计算
- 用Parallel和foreach包玩转并行计算
- Udacity并行计算课程笔记-The GPU Hardware and Parallel Communication Patterns
- 【转】【C#】【Thread】【Parallel】并行计算
- C# Parallel 并行计算测试
- 解决了PyTorch 使用torch.nn.DataParallel 进行多GPU训练的一个BUG:模型(参数)和数据不在相同设备上
- 把多核CPU的计算能力都用起来吧,Parallel--让你的循环变成并行运算
- Parallel并行计算合计数据时错误的原因和解决办法
- linux命令parallel实现多进程并行计算
- 【并行计算】基于OpenMP的并行编程(#pragma omp parallel for)
- opencv 并行计算函数 parallel_for_的使用