本教程中,我们将会学习如何使用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())
输出结果为:
当你没有或只有一个GPU时,结果就会像上图,但如果你有多个GPU时: