pytorch利用多个GPU并行计算多gpu

一时失言乱红尘 2023-06-05 05:27 141阅读 0赞

版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。  
本文链接：https://blog.csdn.net/Answer3664/article/details/98992409  
参考：

https://pytorch.org/docs/stable/nn.html

https://github.com/apachecn/pytorch-doc-zh/blob/master/docs/1.0/blitz\_data\_parallel\_tutorial.md

一、 torch.nn.DataParallel  
torch.nn.DataParallel(module, device\_ids=None, output\_device=None, dim=0)

在正向传递中，模块在每个设备上复制，每个副本处理一部分输入。在向后传递期间，来自每个副本的渐变被加到原始模块中。

module：需要并行处理的模型  
device\_ids：并行处理的设备，默认使用所有的cuda  
output\_device：输出的位置，默认输出到cuda:0  
例子：

>>> net = torch.nn.DataParallel(model, device\_ids=\[0, 1, 2\])  
>>> output = net(input\_var) \# input\_var can be on any device, including CPU  
torch.nn.DataParallel()返回一个新的模型，能够将输入数据自动分配到所使用的GPU上。所以输入数据的数量应该大于所使用的设备的数量。

二、一个完整例子  
import torch  
import torch.nn as nn  
from torch.utils.data import Dataset, DataLoader  
  
\# parameters and DataLoaders  
input\_size = 5  
output\_size = 2  
  
batch\_size = 30  
data\_size = 100  
  
device = torch.device('cuda:0' if torch.cuda.is\_available() else 'cpu')  
  
  
\# 随机数据集  
class RandomDataset(Dataset):  
  
def \_\_init\_\_(self, size, length):  
self.len = length  
self.data = torch.randn(length, size)  
  
def \_\_getitem\_\_(self, index):  
return self.data\[index\]  
  
def \_\_len\_\_(self):  
return self.len  
  
  
rand\_loader = DataLoader(dataset=RandomDataset(input\_size, data\_size),  
batch\_size=batch\_size, shuffle=True)  
  
  
\# 以简单模型为例，同样可以用于CNN, RNN 等复杂模型  
class Model(nn.Module):  
def \_\_init\_\_(self, input\_size, output\_size):  
super(Model, self).\_\_init\_\_()  
self.fc = nn.Linear(input\_size, output\_size)  
  
def forward(self, input):  
output = self.fc(input)  
print('In model: input size', input.size(), 'output size:', output.size())  
return output  
  
  
\# 实例  
model = Model(input\_size, output\_size)  
  
if torch.cuda.device\_count() > 1:  
print("Use", torch.cuda.device\_count(), 'gpus')  
model = nn.DataParallel(model)  
  
model.to(device)  
  
for data in rand\_loader:  
input = data.to(device)  
output = model(input)  
print('Outside: input size ', input.size(), 'output size: ', output.size())  
输出：

In model: input size torch.Size(\[30, 5\]) output size: torch.Size(\[30, 2\])  
Outside: input size  torch.Size(\[30, 5\]) output size:  torch.Size(\[30, 2\])  
In model: input size torch.Size(\[30, 5\]) output size: torch.Size(\[30, 2\])  
Outside: input size  torch.Size(\[30, 5\]) output size:  torch.Size(\[30, 2\])  
In model: input size torch.Size(\[30, 5\]) output size: torch.Size(\[30, 2\])  
Outside: input size  torch.Size(\[30, 5\]) output size:  torch.Size(\[30, 2\])  
In model: input size torch.Size(\[10, 5\]) output size: torch.Size(\[10, 2\])  
Outside: input size  torch.Size(\[10, 5\]) output size:  torch.Size(\[10, 2\])

若有2个GPU

Use 2 GPUs!  
In Model： input size torch.Size(\[15, 5\]) output size torch.Size(\[15, 2\])  
In Model： input size torch.Size(\[15, 5\]) output size torch.Size(\[15, 2\])  
Outside： input size torch.Size(\[30, 5\]) output\_size torch.Size(\[30, 2\])  
In Model： input size torch.Size(\[15, 5\]) output size torch.Size(\[15, 2\])  
In Model： input size torch.Size(\[15, 5\]) output size torch.Size(\[15, 2\])  
Outside： input size torch.Size(\[30, 5\]) output\_size torch.Size(\[30, 2\])  
In Model： input size torch.Size(\[15, 5\]) output size torch.Size(\[15, 2\])  
In Model： input size torch.Size(\[15, 5\]) output size torch.Size(\[15, 2\])  
Outside： input size torch.Size(\[30, 5\]) output\_size torch.Size(\[30, 2\])  
In Model： input size torch.Size(\[5, 5\]) output size torch.Size(\[5, 2\])  
In Model： input size torch.Size(\[5, 5\]) output size torch.Size(\[5, 2\])  
Outside： input size torch.Size(\[10, 5\]) output\_size torch.Size(\[10, 2\])  
若有3个GPU

Use 3 GPUs!  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
Outside： input size torch.Size(\[30, 5\]) output\_size torch.Size(\[30, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
Outside： input size torch.Size(\[30, 5\]) output\_size torch.Size(\[30, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
In Model： input size torch.Size(\[10, 5\]) output size torch.Size(\[10, 2\])  
Outside： input size torch.Size(\[30, 5\]) output\_size torch.Size(\[30, 2\])  
In Model： input size torch.Size(\[4, 5\]) output size torch.Size(\[4, 2\])  
In Model： input size torch.Size(\[4, 5\]) output size torch.Size(\[4, 2\])  
In Model： input size torch.Size(\[2, 5\]) output size torch.Size(\[2, 2\])  
Outside： input size torch.Size(\[10, 5\]) output\_size torch.Size(\[10, 2\])  
总结：

DataParallel自动的划分数据，并将作业发送到多个GPU上的多个模型。DataParallel会在每个模型完成作业后，收集与合并结果然后返回给你。  
————————————————  
版权声明：本文为CSDN博主「Answerlzd」的原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接及本声明。  
原文链接：https://blog.csdn.net/Answer3664/article/details/98992409

![1432127-20190905160448359-185302602.jpg][]

转载于:https://www.cnblogs.com/jfdwd/p/11466683.html

[1432127-20190905160448359-185302602.jpg]: /images/20230601/b3071363165543a482f5e52aca2c3880.png