主要介绍在Pytorch中如何使用GPU进行加速训练，同时在多GPU的条件下如何进行单机并行计算。
一、Pytorch使用GPU加速训练
常用命令
1、 获取设备里的GPU
#syntax 1
device = torch.device("cuda:0,1")
#syntax 2 会根据设备情况自动选择cpu或者gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
2、将模型和数据放到GPU中运行
# syntax 1
model.to(device)
# syntax 2 可以直接使用
model.gpu()
model.cuda()
# 将数据放入gpu
data = data.to(device)
data = data.gpu()
data = data.cuda()
通过以上例子知道如何将模型以及数据放入CPU中加速训练。其他的步骤基本保持不变。
当然，如果想使用cpu训练
device = torch.cuda("cpu")
model.to(device)
data.to(device)
# 或者
model.cpu()
data = data.cpu()
3、
在多 GPU 中执行前馈，后馈操作是非常自然的。尽管如此，PyTorch 默认只会使用一个 GPU。通过使用 DataParallel 让你的模型并行运行，你可以很容易的在多 GPU 上运行你的操作。
首先指定GPU运行的几种方式
# 终端指定
CUDA_VISIBLE_DEVICES=1 python script.py
# python 代码中指定
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"
# 不使用的话就可以指定 "none"，结合前面的device比较方便
然后将模型放放在多GPU运行
model = nn.DataParallel(model)
二、代码实例
引入 PyTorch 模块和定义参数
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
# Parameters and DataLoaders
input_size = 5
output_size = 2
batch_size = 30
data_size = 100
设备
import os
# 这里可以指定设备的id, 或者 "none"
os.environ["CUDA_VISIBLE_DEVICES"] ="0,1"
DEVICE= torch.device("cuda" if torch.cuda.is_available() else "cpu")
# DEVICE = torch.device("cpu")
生成数据
通过使用Dataset和DataLoader对象
class RandomDataset(Dataset):
def __init__(self,size,length):
self.len = length
self.data = torch.randn(length,size)
def __getitem__(self, index: int):
return self.data[index]
def __len__(self):
return self.len
rand_dataset = RandomDataset(input_size,data_size)
rand_loader = DataLoader(rand_dataset,batch_size=batch_size,shuffle=True)
简单模型
为了做一个小 demo，我们的模型只是获得一个输入，执行一个线性操作，然后给一个输出。尽管如此，你可以使用 DataParallel 在任何模型(CNN, RNN, Capsule Net 等等.)
我们放置了一个输出声明在模型中来检测输出和输入张量的大小。
class MyModel(nn.Module):
def __init__(self,input_size,output_size):
super(MyModel,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
创建模型并且数据并行处理
这是整个教程的核心。首先我们需要一个模型的实例，然后验证我们是否有多个 GPU。如果我们有多个 GPU，我们可以用 nn.DataParallel 来 包裹 我们的模型。然后我们使用 model.to(device) 把模型放到多 GPU 中。
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
if torch.cuda.is_available():
model.cuda()
＃ 前面定义了device，所以可以
model.to(device)
运行程序
for data in rand_loader:
if torch.cuda.is_available():
input_var = Variable(data.cuda())
else:
input_var = Variable(data)
# 或者 input_var = data.to(device)
output = model(input_var)
print("Outside: input size", input_var.size(),
"output_size", output.size())
结果
使用1个GPU
使用1个GPU
使用2个GPU
使用2个GPU
使用3个GPU
使用3个GPU
以上就学会了使用GPU加速模型训练，同时知道如何使用多GPU并行化训练。
网友评论