作者|DR. VAIBHAV KUMAR
編譯|VK
來源|Analytics In Diamag
隨著深度學習模型在各種應用中的成功實施，現在是時候獲得不僅準確而且速度更快的結果，
為了得到更準確的結果，資料的大小是非常重要的，但是當這個大小影響到機器學習模型的訓練時間時，這一直是一個值得關注的問題，
為了克服訓練時間的問題，我們使用TPU運行時環境來加速訓練，為此，PyTorch一直在通過提供最先進的硬體加速器來支持機器學習的實作，
PyTorch對云TPU的支持是通過與XLA（加速線性代數）的集成實作的，XLA是一種用于線性代數的編譯器，可以針對多種型別的硬體，包括CPU、GPU和TPU，
本文演示了如何使用PyTorch和TPU實作深度學習模型，以加快訓練程序，
在這里，我們使用PyTorch定義了一個卷積神經網路（CNN）模型，并在PyTorch/XLA環境中對該模型進行了訓練，
XLA將CNN模型與分布式多處理環境中的Google Cloud TPU（張量處理單元）連接起來，在這個實作中，使用8個TPU核心來創建一個多處理環境，
我們將用這個PyTorch深度學習框架進行時裝分類測驗，觀察訓練時間和準確性，
用PyTorch和TPU實作CNN
我們將在Google Colab中實作執行，因為它提供免費的云TPU（張量處理單元），
在繼續下一步之前，在Colab筆記本中，轉到“編輯”，然后選擇“設定”，從下面螢屏截圖中的串列中選擇“TPU”作為“硬體加速器”，
驗證TPU下面的代碼是否正常運行，
import os
assert os.environ['COLAB_TPU_ADDR']
如果啟用了TPU，它將成功執行，否則它將回傳‘KeyError: ‘COLAB_TPU_ADDR’’，你也可以通過列印TPU地址來檢查TPU，
TPU_Path = 'grpc://'+os.environ['COLAB_TPU_ADDR']
print('TPU Address:', TPU_Path)
啟用TPU后，我們將安裝兼容的控制盤和依賴項，以使用以下代碼設定XLA環境，
VERSION = "20200516"
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version $VERSION
一旦安裝成功，我們將繼續定義加載資料集、初始化CNN模型、訓練和測驗的方法，首先，我們將匯入所需的庫，
import numpy as np
import os
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.debug.metrics as met
import torch_xla.distributed.parallel_loader as pl
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.utils.utils as xu
from torchvision import datasets, transforms
之后，我們將進一步定義需要的超引數，
# 定義引數
FLAGS = {}
FLAGS['datadir'] = "/tmp/mnist"
FLAGS['batch_size'] = 128
FLAGS['num_workers'] = 4
FLAGS['learning_rate'] = 0.01
FLAGS['momentum'] = 0.5
FLAGS['num_epochs'] = 50
FLAGS['num_cores'] = 8
FLAGS['log_steps'] = 20
FLAGS['metrics_debug'] = False
下面的代碼片段將把CNN模型定義為PyTorch實體，以及用于加載資料、訓練模型和測驗模型的函式，
SERIAL_EXEC = xmp.MpSerialExecutor()
class FashionMNIST(nn.Module):
def __init__(self):
super(FashionMNIST, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.bn1 = nn.BatchNorm2d(10)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.bn2 = nn.BatchNorm2d(20)
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = self.bn1(x)
x = F.relu(F.max_pool2d(self.conv2(x), 2))
x = self.bn2(x)
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)
# 只在記憶體中實體化一次模型權重，
WRAPPED_MODEL = xmp.MpModelWrapper(FashionMNIST())
def train_mnist():
torch.manual_seed(1)
def get_dataset():
norm = transforms.Normalize((0.1307,), (0.3081,))
train_dataset = datasets.FashionMNIST(
FLAGS['datadir'],
train=True,
download=True,
transform=transforms.Compose(
[transforms.ToTensor(), norm]))
test_dataset = datasets.FashionMNIST(
FLAGS['datadir'],
train=False,
download=True,
transform=transforms.Compose(
[transforms.ToTensor(), norm]))
return train_dataset, test_dataset
#使用串行執行器可以避免多個行程下載相同的資料
train_dataset, test_dataset = SERIAL_EXEC.run(get_dataset)
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset,
num_replicas=xm.xrt_world_size(),
rank=xm.get_ordinal(),
shuffle=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=FLAGS['batch_size'],
sampler=train_sampler,
num_workers=FLAGS['num_workers'],
drop_last=True)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=FLAGS['batch_size'],
shuffle=False,
num_workers=FLAGS['num_workers'],
drop_last=True)
# 調整學習率
lr = FLAGS['learning_rate'] * xm.xrt_world_size()
# 獲取損失函式、優化器和模型
device = xm.xla_device()
model = WRAPPED_MODEL.to(device)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=FLAGS['momentum'])
loss_fn = nn.NLLLoss()
def train_fun(loader):
tracker = xm.RateTracker()
model.train()
for x, (data, target) in enumerate(loader):
optimizer.zero_grad()
output = model(data)
loss = loss_fn(output, target)
loss.backward()
xm.optimizer_step(optimizer)
tracker.add(FLAGS['batch_size'])
if x % FLAGS['log_steps'] == 0:
print('[xla:{}]({}) Loss={:.5f}'.format(
xm.get_ordinal(), x, loss.item(), time.asctime()), flush=True)
def test_fun(loader):
total_samples = 0
correct = 0
model.eval()
data, pred, target = None, None, None
for data, target in loader:
output = model(data)
pred = output.max(1, keepdim=True)[1]
correct += pred.eq(target.view_as(pred)).sum().item()
total_samples += data.size()[0]
accuracy = 100.0 * correct / total_samples
print('[xla:{}] Accuracy={:.2f}%'.format(
xm.get_ordinal(), accuracy), flush=True)
return accuracy, data, pred, target
# 訓練和評估回圈
accuracy = 0.0
data, pred, target = None, None, None
for epoch in range(1, FLAGS['num_epochs'] + 1):
para_loader = pl.ParallelLoader(train_loader, [device])
train_fun(para_loader.per_device_loader(device))
xm.master_print("Finished training epoch {}".format(epoch))
para_loader = pl.ParallelLoader(test_loader, [device])
accuracy, data, pred, target = test_fun(para_loader.per_device_loader(device))
if FLAGS['metrics_debug']:
xm.master_print(met.metrics_report(), flush=True)
return accuracy, data, pred, target
現在，要將結果繪制為測驗影像的預測標簽和實際標簽，將使用以下功能模塊，
# 結果可視化
import math
from matplotlib import pyplot as plt
M, N = 5, 5
RESULT_IMG_PATH = '/tmp/test_result.png'
def plot_results(images, labels, preds):
images, labels, preds = images[:M*N], labels[:M*N], preds[:M*N]
inv_norm = transforms.Normalize((-0.1307/0.3081,), (1/0.3081,))
num_images = images.shape[0]
fig, axes = plt.subplots(M, N, figsize=(12, 12))
fig.suptitle('Predicted Lables')
for i, ax in enumerate(fig.axes):
ax.axis('off')
if i >= num_images:
continue
img, label, prediction = images[i], labels[i], preds[i]
img = inv_norm(img)
img = img.squeeze() # [1,Y,X] -> [Y,X]
label, prediction = label.item(), prediction.item()
if label == prediction:
ax.set_title(u'Actual {}/ Predicted {}'.format(label, prediction), color='blue')
else:
ax.set_title(
'Actual {}/ Predicted {}'.format(label, prediction), color='red')
ax.imshow(img)
plt.savefig(RESULT_IMG_PATH, transparent=True)
現在，我們都準備好在MNIST資料集上訓練模型，訓練開始前，我們將記錄開始時間，訓練結束后，我們將記錄結束時間并列印50個epoch的總訓練時間，
# 啟動訓練流程
def train_cnn(rank, flags):
global FLAGS
FLAGS = flags
torch.set_default_tensor_type('torch.FloatTensor')
accuracy, data, pred, target = train_mnist()
if rank == 0:
# 檢索TPU核心0上的張量并繪制，
plot_results(data.cpu(), pred.cpu(), target.cpu())
xmp.spawn(train_cnn, args=(FLAGS,), nprocs=FLAGS['num_cores'],
start_method='fork')
一旦訓練成功結束，我們將列印訓練所用的總時間，
end_time = time.time()
print('Total Training time = ',end_time-start_time )
正如我們在上面看到的，這種方法花費了269秒或大約4.5分鐘，這意味著50個epoch訓練PyTorch模型不到5分鐘，最后，我們將通過訓練的模型來可視化預測，
from google.colab.patches import cv2_imshow
import cv2
img = cv2.imread(RESULT_IMG_PATH, cv2.IMREAD_UNCHANGED)
cv2_imshow(img)
因此，我們可以得出這樣的結論：使用TPU實作深度學習模型可以實作快速的訓練，正如我們前面所看到的那樣，
在不到5分鐘的時間內，對50個epoch的40000張訓練影像進行了CNN模型的訓練，我們在訓練中也獲得了89%以上的準確率，
因此，在TPU上訓練深度學習模型在時間和準確性方面總是有好處的，
參考文獻：
- Joe Spisak, “Get started with PyTorch, Cloud TPUs, and Colab”.
- “PyTorch on XLA Devices”, PyTorch release.
- “Training PyTorch models on Cloud TPU Pods”, Google Cloud Guides.
原文鏈接：https://analyticsindiamag.com/how-to-implement-cnn-model-using-pytorch-with-tpu/
歡迎關注磐創AI博客站：
http://panchuang.net/
sklearn機器學習中文官方檔案：
http://sklearn123.com/
歡迎關注磐創博客資源匯總站：
http://docs.panchuang.net/
轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/2254.html
標籤：其他
上一篇：C++機器學習庫介紹
下一篇：構建物件檢測模型