在神经网络领域，许多人可能使用CPU来训练模型。对于只有十万参数的模型来说，即使训练需要几个小时，这也是可以接受的。但如果模型参数达到十亿或二十亿，比如常见的CNN模型VGG16有1.38亿参数，那么使用CPU训练就会变得不切实际，因为它会耗费大量的宝贵时间。本文将讨论GPU如何帮助解决这一问题，并提供使用GPU训练一个简单模型的实践经验。
与其让介绍GPU的优势，不如相信这个惊人的视频会给更好的理解。视频展示了GPU计算的强大能力。现在可能已经有所了解，GPU的大规模并行计算能力极大地帮助提高了性能并减少了复杂神经网络模型的训练时间。GPU包含大量内置的小型核心，这在执行任务时非常有帮助。
神经网络中最基础的操作是矩阵乘法，GPU在这方面非常擅长，它解决这些计算问题就像一个专业的数学家专门从事矩阵乘法一样。GPU相对于CPU的其他优势包括：
为了有效使用GPU的多个核心，使用CUDA编程模型。在PyTorch中运行CUDA操作非常容易。但请记住，GPU不会完全取代CPU，因为GPU只会作为CPU的额外贡献者，帮助运行给定应用程序的并行重复计算，而主程序仍然在CPU上运行。GPU比CPU更适用的其他应用包括：
与其过多地停留在理论层面，不如让开始在Google Colab笔记本上使用GPU训练模型。
使用Google Colab环境，可以免费访问“NVIDIA Tesla K80”GPU。但请注意，只能连续使用它12小时，之后可能在特定时间内无法访问它，除非购买Colab Pro。将以MNIST手写数字分类数据集作为问题。任务是训练一个模型，它可以正确地将给定的手写数字图像分类到相应的标签。因此，在GPU上训练模型时，需要强调的一些主要步骤包括：
让开始吧！
在创建一个新的笔记本后，第一步是将运行时类型设置为GPU。
import torch
import torchvision
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torchvision.utils import make_grid
from torch.utils.data.dataloader import DataLoader
from torch.utils.data import random_split
%matplotlib inline
# Use a white background for matplotlib figures
matplotlib.rcParams['figure.facecolor'] = '#ffffff'
dataset = MNIST(root='data/', download=True, transform=ToTensor())
val_size = 10000
train_size = len(dataset) - val_size
train_set, val_set = random_split(dataset, [train_size, val_size])
batch_size = 128
train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2, pin_memory=True)
val_loader = DataLoader(val_set, batch_size=256, shuffle=True, num_workers=2, pin_memory=True)
当想要将已经加载到CPU上的数据集推送到GPU时，设置pin_memory=True可以加快两者之间的数据传输。
class Mnistmodel(nn.Module):
def __init__(self, input_size, hidden_size, out_size):
super().__init__()
self.linear1 = nn.Linear(input_size, hidden_size)
self.linear2 = nn.Linear(hidden_size, out_size)
def forward(self, xb):
xb = xb.view(xb.size(0), -1)
out = self.linear1(xb)
rel = F.relu(out)
out = self.linear2(rel)
return out
def training_step(self, batch):
images, labels = batch
out = self(images)
loss = F.cross_entropy(out, labels)
return loss
def validation_step(self, batch):
images, labels = batch
out = self(images)
loss = F.cross_entropy(out, labels)
acc = accuracy(out, labels)
return {'val_loss': loss, 'val_acc': acc}
def validation_epoch_end(self, outputs):
batch_losses = [x['val_loss'] for x in outputs]
epoch_loss = torch.stack(batch_losses).mean()
batch_accuracy = [x['val_acc'] for x in outputs]
epoch_accuracy = torch.stack(batch_accuracy).mean()
return {'val_loss': epoch_loss, 'val_acc': epoch_accuracy}
def epoch_end(self, epoch, result):
print("Epoch [{}], val_loss: {:.4f}, val_acc: {:.4f}".format(epoch, result['val_loss'], result['val_acc']))
使用这个类，可以创建需要训练的模型。在将模型和数据加载到GPU之前，让检查一下GPU是否可用？
def get_default_device():
if torch.cuda.is_available():
return torch.device('cuda')
else:
return torch.device('cpu')
现在即使GPU不可用，也不会有问题，因为当GPU不可用时，将自动切换到CPU进行训练，但训练时间会很长。
def to_device(data, device):
if isinstance(data, (list, tuple)):
return [to_device(x, device) for x in data]
return data.to(device, non_blocking=True)
class DeviceDataLoader():
def __init__(self, dl, device):
self.dl = dl
self.device = device
def __iter__(self):
for b in self.dl:
yield to_device(b, self.device)
def __len__(self):
return len(self.dl)
使用DevicedataLoader类，将创建对象，帮助将train_loader和val_loader中的数据（在步骤4中定义）移动到GPU。
def fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):
history = []
optimizer = opt_func(model.parameters(), lr)
for epoch in range(epochs):
for batch in train_loader:
loss = model.training_step(batch)
loss.backward()
optimizer.step()
optimizer.zero_grad()
result = evaluate(model, val_loader)
model.epoch_end(epoch, result)
history.append(result)
return history
def evaluate(model, val_loader):
outputs = [model.validation_step(batch) for batch in val_loader]
return model.validation_epoch_end(outputs)
num_classes = 10
model = Mnistmodel(input_size=784, hidden_size=32, out_size=num_classes)
to_device(model, device)
请注意，在创建模型后，必须将其移动到GPU，否则数据已经在GPU中，而模型还在CPU中。