data parallelism

本教程中,我们将会学习如何使用DataParallel来使用多个GPU。

使用PyTorch来使用GPU非常容易,你可以直接把模型放到GPU上:

## 这里是为了举例,正常使用前应检查cuda可用性
device = torch.device("cuda:0")
net.to(device)

然后,你可以将所有的张量转移到GPU上:

mytensor = my_tensor.to(device)

注意,只是调用my_tensor(device)的话,是返回一个my_tensor在GPU上的复制,而不是重写my_tensor。你需要把它指派给一个新张量并在GPU上使用这个新张量。

在多个GPU上执行前向和后向传播是很自然的,但是PyTorch默认只使用一个GPU,你可以简单地通过使用DataParallel来让模型在多个GPU上并行运行。

model = nn.DataParallel(model)

引用和参数

引用PyTorch组件和定义参数

import torch
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader

input_size = 5
output_size = 2

batch_size = 30
data_size = 100

定义设备

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

假数据集

设计一个假的(随机)数据集,你只需应用getitem

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 = torch.utils.data.DataLoader(dataset=RandomDataset(input_size,data_size),
                     batch_size=batch_size,shuffle=True)

简单模型

为了演示,我们的模型只获取输入,执行线性操作并提供输出。然后,你可以将DataParallel应用于任何模型上(CNN,RNN,Capsule Net等)。

我们在模型中放置了一个输出语句来监视输入和输出张量的大小。请注意0级批次上打印的内容。

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("\t模型的输入大小:",input.size(),"输出大小:",output.size())
    
   	 	return output

创建模型和数据并行

这是本教程的核心部分。首先,我们需要创建一个模型实例,并检查我们是否有多个GPU。如果我们有多个GPU,我们可以使用nn.DataParallel来包裹我们的模型。然后我们可以使用 model.to(device) 将其放入GPU中。

model = Model(input_size,output_size)
if torch.cuda.device_count()> 1:
	print("Let's use %d GPUs" % torch.cuda.device_count())
	model = nn.DataParallel(model)

model.to(device)

运行模型

现在我们可以看看输入和输出张量的大小。

for data in rand_loader:
	input = data.to(device)
	output = model(input)
    print("\t输入张量大小:",input.size(),"输出张量大小:",output.size())

输出结果为:

data parallelism
当你没有或只有一个GPU时,结果就会像上图,但如果你有多个GPU时:

data parallelism
data parallelism
data parallelism

上一篇:Flink Parallelism 与 Slot 理解


下一篇:在Ubuntu上并行运行作业 – Perl和Java之间的I / O争用差异