![](https://i-blog.csdnimg.cn/blog_migrate/4aa08fa61c52729a178d7d83a032df70.png)
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. 创建虚拟数据集
为了简单起见,创建一个简单的虚拟数据集,由随机值张量生成。
# 创建虚拟数据集类
class RandomDataset(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运算应用到任何复杂的神经网络模型中。
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("\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: # 检查电脑是否有多块GPU
print(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运算了。