PyTorch運算加速

簡介

在前一篇文章中完整演示了整個PyTorch進行深度模型訓練的流程，然而，盡管大多時候就是按照這個思路進行模型訓練的，但是我們實際上忽略了一個至關重要的問題，那就是GPU加速，目前主流的深度學習框架都是支持GPU加速運算的，PyTorch也不例外，本文介紹如何使用GPU加速模型的訓練。

GPU

事實上，想要使用GPU加速深度學習模型的訓練就是將相關的內存中的變量轉移到顯存中，利用顯卡遠超CPU的運算速度得到結果。

為了將這些數據搬到顯存上，需要借助CUDA這個運算平臺（所以機器上必須安裝合適的CUDA和Cudnn），然后將內存變量通過.cuda()方法轉移到GPU上，這個方法由PyTorch封裝完成，訓練過程中由于顯存的限制一般是將模型以及當前批次的數據轉移到顯存中。

本文在T4顯卡上進行演示，具體配置如下。

通過下面的代碼判斷PyTorch是否可以使用GPU加速，以及可用的GPU數目。

import torch print(torch.__version__) print("gpu can be used", torch.cuda.is_available()) print("gpu device", torch.cuda.device_count())

例如本機環境的測試結果如下。

GPU加速訓練

這一步主要的操作就是模型、每一批的數據搬到GPU上，預測的結果由于用于可視化工具或者命令行輸出，所以需要是CPU類型，通過.cpu()方法將顯存上的數據拿到內存中。主要的代碼如下，涉及到修改的地方都用注釋說明了。下面的案例同之前的文章。

import pandas as pd import torch from torch.utils.data import Dataset, DataLoader from torchvision.transforms import transforms import torch.nn as nn import torch.nn.functional as F from torch import optim from PIL import Image import timestart = time.time() print("start time", start)class MyDataset(Dataset):def __init__(self, desc_file, transform=None):self.all_data = pd.read_csv(desc_file).valuesself.transform = transformdef __getitem__(self, index):img, label = self.all_data[index, 0], self.all_data[index, 1]img = Image.open(img).convert('RGB')if self.transform is not None:img = self.transform(img)return img, labeldef __len__(self):return len(self.all_data)class Net(nn.Module):def __init__(self):super(Net, self).__init__()self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=(3, 3))self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)self.pool2 = nn.MaxPool2d(2, 2)self.fc1 = nn.Linear(64*6*6, 256)self.fc2 = nn.Linear(256, 128)self.fc3 = nn.Linear(128, 101)def forward(self, x):x = self.pool1(F.relu(self.conv1(x)))x = self.pool2(F.relu(self.conv2(x)))x = x.view(-1, 64*6*6)x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return xdef init_weights(self):for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.xavier_normal_(m.weight.data)if m.bias is not None:m.bias.data.zero_()elif isinstance(m, nn.BatchNorm2d):m.weight.data.fill_(1)m.bias.data.zero_()elif isinstance(m, nn.Linear):nn.init.normal_(m.weight.data, 0, 0.01)m.bias.data.zero_()desc_train = '../data/desc_train.csv' desc_valid = '../data/desc_valid.csv' desc_test = '../data/desc_test.csv'batch_size = 32 lr = 0.001 epochs = 50 norm_mean = [0.4948052, 0.48568845, 0.44682974] norm_std = [0.24580306, 0.24236229, 0.2603115]train_transform = transforms.Compose([transforms.Resize((32, 32)),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize(norm_mean, norm_std) # 按照imagenet標準 ])valid_transform = transforms.Compose([transforms.Resize((32, 32)),transforms.ToTensor(),transforms.Normalize(norm_mean, norm_std) ])test_transform = transforms.Compose([transforms.Resize((32, 32)),transforms.ToTensor(),transforms.Normalize(norm_mean, norm_std) ])train_data = MyDataset(desc_train, transform=train_transform) valid_data = MyDataset(desc_valid, transform=valid_transform) test_data = MyDataset(desc_test, transform=test_transform)# 構建DataLoader train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True) valid_loader = DataLoader(dataset=valid_data, batch_size=batch_size) test_loader = DataLoader(dataset=test_data, batch_size=batch_size)model = Net() model.init_weights() # !!! change here !!! model.cuda()criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, dampening=0.1) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, verbose=True)for epoch in range(epochs):# 訓練集訓練train_loss = 0.0correct = 0.0total = 0.0for step, data in enumerate(train_loader):x, y = data# !!! change here !!!x = x.cuda()y = y.cuda()out = model(x)loss = criterion(out, y)optimizer.zero_grad()loss.backward()optimizer.step()_, pred = torch.max(out.data, 1)total += y.size(0)# !!! change here !!!correct += (pred == y).squeeze().sum().cpu().numpy()train_loss += loss.item()if step % 100 == 0:print("epoch", epoch, "step", step, "loss", loss.item())train_acc = correct / total# 驗證集驗證valid_loss = 0.0correct = 0.0total = 0.0for step, data in enumerate(valid_loader):model.eval()x, y = data# !!! change here !!!x = x.cuda()y = y.cuda()out = model(x)out.detach_()loss = criterion(out, y)_, pred = torch.max(out.data, 1)valid_loss += loss.item()total += y.size(0)# !!! change here !!!correct += (pred == y).squeeze().sum().cpu().numpy()valid_acc = correct / totalscheduler.step(valid_loss)print("finish training")net_save_path = 'net_params.pkl' torch.save(model.state_dict(), net_save_path)end = time.time() print("end time", end) print("spend time", end-start)

為了對比GPU加速的效果，下面同樣的訓練配置（輪數、數據等）下，對比CPU和GPU下的花費訓練時間。注意，一開始GPU需要warm up，所以一開始的速度不會很快，之后會恢復正常。

下圖是CPU和GPU花費的訓練時間對比，顯然，GPU快了很多，事實上，這里數據較少，采用大量數據會更加凸顯GPU加速的優勢。

補充說明

本文主要演示了PyTorch如何使用GPU加速模型的訓練，事實上只要將模型和數據遷移到顯存上運算即可。本文涉及到的所有代碼均可以在我的Github找到，歡迎star或者fork。

超強干貨來襲 云風專訪：近40年碼齡，通宵達旦的技術人生

總結

以上是生活随笔為你收集整理的PyTorch-运算加速的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。