【PyTorch教程】07-PyTorch如何使用多块GPU训练神经网络模型

本期目录

1. 绪论

在本篇博文中，你将学习到在PyTorch中如何使用多GPU进行并行训练。在PyTorch中使用多GPU训练神经网络模型是非常简单的，PyTorch已经为我们封装好一个

nn.DataParallel

类来进行多GPU训练。

先来回顾一下在单GPU中如何训练，首先，我们可以把张量复制到GPU上：

import torch
import torch.nn as nn

device = torch.device("cuda:0")# 创建单GPU对象
mytensor = my_tensor.to(device)# 复制到该GPU上

我们还可以把神经网络模型放到GPU上：

model.to(device)

【注意】

原本在CPU的张量只是复制了一个副本过去GPU，其本身仍旧存储在CPU上。因此使用

torch.to('cuda')

函数时，要声明一个新的变量来接收GPU张量。

但是，PyTorch默认只会使用一张GPU。因此我们可以使用

nn.DataParallel

在多块GPU上并行地运行你的神经网络模型：

model = nn.DataParallel(model)

下面将是本篇教程的核心内容。我会用一个简单的数据集和神经网络模型来详细演示如何使用多GPU进行运算。

2. 导入Pytoch模块并声明参数

为了简单起见，我们声明一个输入尺寸为5，输出尺寸为2的张量。批量大小设置为30，数据集大小为100：

mport 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

声明GPU：

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

3. 创建虚拟数据集

为了简单起见，创建一个简单的虚拟数据集，由随机值张量生成。

# 创建虚拟数据集类classRandomDataset(Dataset):# 继承torch.utils.data.Dataset抽象父类，要实现getitem抽象方法def__init__(self, size, length):# RandomDataset类的构造器
        self.len= length
        self.data = torch.randn(length, size)# 创建一个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, num_workers=12)

4. 搭建一个简单的模型

简单起见，搭建一个只有一层全连接层的神经网络模型。为了方便大家观察到输入输出变化，我在模型中插入了

print

语句，在实际使用中不需要加。当然，你可以把本文的多GPU运算应用到任何复杂的神经网络模型中。

classModel(nn.Module):def__init__(self, input_size, output_size):super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)defforward(self,input):
        output = self.fc(input)print("\tIn Model: input size",input.size,"output size", output.size)return output

5. 多GPU并行计算

这部分内容是本文的核心内容。首先，我们实例化模型对象，然后检查我们的电脑是否有多块GPU。如果有多块GPU，就可以使用

nn.DataParallel

打包我们的模型，最后用

nn.DataParallel

的

model.to(device)

把模型传送到多块GPU中进行运算。

model = Model(input_size, output_size)# 实例化模型对象if torch.cuda.device_count()>1:# 检查电脑是否有多块GPUprint(f"Let's use {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)# 将模型对象转变为多GPU并行运算的模型

model.to(device)# 把并行的模型移动到GPU上

6. 运行模型

现在，我们可以让模型跑起来，同时大家可以观察一下输入输出张量的形状变化：

for data in rand_loader:input= data.to(device)# 把输入张量移到GPU
    output = model(input)print(f"Outside: input size {input.size()},"f"output_size {output.size()}")

博主服务器里有8张GPU，所以可以看到如下输出：

Let's use 8 GPUs!
    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([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([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([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([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([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([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([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([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([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([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([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([30, 5]) output_size torch.Size([30, 2])
    In Model: input size torch.Size([2, 5])    output size torch.Size([2, 2])
    In Model: input size torch.Size([2, 5])    output size torch.Size([2, 2])
    In Model: input size torch.Size([2, 5])    output size torch.Size([2, 2])
    In Model: input size torch.Size([2, 5])    output size torch.Size([2, 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])

可以看到，我们一共有100个数据，每个批量30个。PyTorch会自动帮我们尽可能地把每批次数据平均分到8张卡上。因此在前3个批次中，前7张卡都分到了4个数据，而最后一张卡分剩下的2个数据，一共是

     7
    
    
     ×
    
    
     4
    
    
     +
    
    
     2
    
    
     =
    
    
     30
    
   
   
     7\times 4 + 2=30 
   
  
 7×4+2=30

个数据。刚好等于

batch_size = 30

。

而前3个批次已经计算完

    30
   
   
    ×
   
   
    3
   
   
    =
   
   
    90
   
  
  
   30\times3=90
  
 
30×3=90 个数据了，自然剩下的 

 
  
   
    100
   
   
    −
   
   
    90
   
   
    =
   
   
    10
   
  
  
   100-90=10
  
 
100−90=10 个数据就分到了5张卡上，每张2个数据。100个数据分4个批次就运算完毕了，每个批次8张GPU都是并行地运算。

至此，恭喜你已经掌握了如何使用PyTorch进行多GPU运算了。