PyTorch實現TPU版本CNN模型

2021-01-11 人工智慧遇見磐創

隨著深度學習模型在各種應用中的成功實施，現在是時候獲得不僅準確而且速度更快的結果。

為了得到更準確的結果，數據的大小是非常重要的，但是當這個大小影響到機器學習模型的訓練時間時，這一直是一個值得關注的問題。

為了克服訓練時間的問題，我們使用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 osassert 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 npimport osimport timeimport torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optimimport torch_xlaimport torch_xla.core.xla_model as xmimport torch_xla.debug.metrics as metimport torch_xla.distributed.parallel_loader as plimport torch_xla.distributed.xla_multiprocessing as xmpimport torch_xla.utils.utils as xufrom torchvision import datasets, transforms之後，我們將進一步定義需要的超參數。

# 定義參數FLAGS = {}FLAGS['datadir'] = "/tmp/mnist"FLAGS['batch_size'] = 128FLAGS['num_workers'] = 4FLAGS['learning_rate'] = 0.01FLAGS['momentum'] = 0.5FLAGS['num_epochs'] = 50FLAGS['num_cores'] = 8FLAGS['log_steps'] = 20FLAGS['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 mathfrom matplotlib import pyplot as pltM, N = 5, 5RESULT_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_imshowimport cv2img = 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.

相關焦點

  • PyTorch終於能用上谷歌雲TPU,推理性能提升4倍,該如何薅羊毛?
    現在PyTorch官方已經在Github上給出示例代碼，教你如何免費使用谷歌雲TPU訓練模型，然後在Colab中進行推理。訓練ResNet-50PyTorch先介紹了在雲TPU設備上訓練ResNet-50模型的案例。如果你要用雲TPU訓練其他的圖像分類模型，操作方式也是類似的。
  • 從R-CNN到YOLO,一文帶你了解目標檢測模型(附論文下載)
    https://github.com/hollance/YOLO-CoreML-MPSNNGraph安卓上基於TensorFlow框架運行YOLO模型實現實時目標檢測。https://github.com/natanielruiz/android-yoloYOLOv2模型時隔一年，YOLO作者放出了v2版本，稱為YOLO9000，並直言它「更快、更高、更強」。YOLO v2的主要改進是提高召回率和定位能力。
  • 13個算法工程師必須掌握的PyTorch Tricks
    2、查看模型每層輸出詳情Keras有一個簡潔的API來查看模型的每一層輸出尺寸，這在調試網絡時非常有用。現在在PyTorch中也可以實現這個功能。然後我將Pytorch升級到1.2版本，試用了下 one_hot 函數，確實很方便。
  • 雲計算學習:用PyTorch實現一個簡單的分類器
    回想了一下自己關於 pytorch 的學習路線，一開始找的各種資料，寫下來都能跑，但是卻沒有給自己體會到學習的過程。有的教程一上來就是寫一個 cnn，雖然其實內容很簡單，但是直接上手容易讓人找不到重點，學的雲裡霧裡。有的教程又淺嘗輒止，師傅領到了門檻跟前，總感覺自己還沒有進門，教程就結束了。
  • Facebook 發布 Detectron2:基於 PyTorch 的新一代目標檢測工具
    而 Detectron2 則是剛推出的新版 PyTorch 1.3 （詳情查看：https://www.leiphone.com/news/201910/9J4PTd833WMwdVoQ.html）中一重大新工具，它源於 maskrcnn 基準測試，但它卻是對先前版本 detectron 的一次徹底重寫。
  • 用Java實現目標檢測|PyTorch
    魚羊 編輯整理量子位 報導 | 公眾號 QbitAI編者按：作為一個Java開發者，你是否曾為在PyTorch上部署模型而苦惱？這篇來自AWS軟體工程師的投稿，結合實例，詳細介紹了DJL這個為Java開發者設計的深度學習庫：5分鐘，你就能在PyTorch上，用Java實現目標檢測。
  • 如何在PyTorch和TensorFlow中訓練圖像分類模型
    以下是兩個傑出的代表：TensorBoard：使用數據流圖幫助有效地可視化數據TensorFlow：對於快速部署新算法/實驗非常有用TensorFlow當前正在運行2.0版本，該版本於2019年9月正式發布。我們還將在2.0版本中實現CNN。我希望你現在對PyTorch和TensorFlow都有基本的了解。現在，讓我們嘗試使用這兩個框架構建深度學習模型並了解其內部工作。
  • 用Java實現目標檢測 | PyTorch
    很多論文都選擇使用PyTorch去實現也證明了它在訓練方面的效率以及易用性。在PyTorch領域，儘管部署一個模型有很多選擇，可為Java開發人員準備的選項卻屈指可數。在過去，用戶可以用PyTorch C++ 寫JNI (Java Native Interface) 來實現這個過程。
  • Pytorch 0.3.0 發布,實現多方面提速,增加對ONNX支持
    經過將近一年的發展，日前，迎來了 PyTorch 0.3.0 版本。根據官方文檔的介紹，此次增加了多個函數和功能，多方面的性能均實現提升。重大變更0.3 版本中刪掉了 Variable.reinforce() 等隨機函數，因為它們的功能有限，並且會影響到性能。
  • 視覺工具包torchvision重大更新:支持分割模型、檢測模型
    有了這些算子，搭建目標檢測模型就更加輕鬆了：· roi_pool (以及模塊版本RoIPool)· roi_align (以及模塊版本RoIAlign)· nms，給邊界框做非極大抑制 (Non-Maximum Suppression用的)· box_iou，用來計算兩組邊界框之間的交集· box_area, 用來計算一組邊界框的面積等到下次更新，這些算子就會支持
  • ResNet——CNN經典網絡模型詳解(pytorch實現)
    ResNet的結構可以極快的加速神經網絡的訓練，模型的準確率也有比較大的提升。同時ResNet的推廣性非常好，甚至可以直接用到InceptionNet網絡中。下圖是ResNet34層模型的結構簡圖。代碼註：本次訓練集下載在AlexNet博客有詳細解說:https://blog.csdn.net/weixin_44023658/article/details/105798326使用遷移學習方法實現收錄在我的這篇blog中： 遷移學習 TransferLearning—通俗易懂地介紹（pytorch實例）#model.pyimport torch.nn
  • DenseNet——CNN經典網絡模型詳解(pytorch實現)
    DenseNet——CNN經典網絡模型詳解(pytorch實現)一個CV小白，寫文章目的為了讓和我意義的小白輕鬆如何，讓大佬鞏固基礎（手動狗頭），大家有任何問題可以一起在評論區留言討論~一、概述論文：Densely Connected Convolutional Networks 論文連結：https://arxiv.org/pdf
  • PyTorch最新:分布式模型並行、Java程序、移動端等多項新功能
    本次更新是最後一個支持 Python2 的版本，同時增加了對分布式模型並行、移動端、Java 程序等方面的支持。 值得注意的是，本次 PyTorch 更新是最後一個支持 Python2 的版本，同時也是最後一個支持 C++11 的版本。官方提示說，用戶應當開始遷移到 Python3，並使用 C++14 開始編譯工作。對於其他詳細信息，感興趣的讀者可以查閱更新文檔。
  • PyTorch 0.2發布:更多NumPy特性,高階梯度、分布式訓練等
    這一版本引入了NumPy風格的Broadcasting、Advanced Indexing，增加了高階梯度和分布式PyTorch。高階梯度PyTorch 0.2版本新增的高階梯度計算支持torch.XXX函數和最流行的nnlayers，在下一版本中可能會支持更多類型。為了支持高階梯度，PyTorch 0.2引入了一種編寫函數autograd.Function的新形式，向下兼容舊的函數編寫形式。
  • PyTorch最新:支持分布式模型並行、Java程序、移動端等多項新功能
    本次更新是最後一個支持 Python2 的版本，同時增加了對分布式模型並行、移動端、Java 程序等方面的支持。值得注意的是，本次 PyTorch 更新是最後一個支持 Python2 的版本，同時也是最後一個支持 C++11 的版本。官方提示說，用戶應當開始遷移到 Python3，並使用 C++14 開始編譯工作。對於其他詳細信息，感興趣的讀者可以查閱更新文檔。
  • 新版PyTorch 1.2 已發布:功能更多、兼容更全、操作更快!
    TorchScript 是一種使用 PyTorch 代碼創建可序列化和可優化模型的方法；任何 TorchScript 程序都可以從 Python 進程中保存，並在不依賴於 Python 的進程中實現加載。
  • Mask R-CNN官方實現「又」來了!基於PyTorch,訓練速度是原來2倍
    該項目旨在讓用戶更容易地創建一種模塊，實現對圖片中物品的識別與分割。負責這一項目的Facebook AI工程師Francisco Massa還在個人Twitter發了張自拍照，並使用新工具對圖片中的物品進行了識別，尷尬的是好像有些對象被識別錯誤：
  • PyTorch中使用DistributedDataParallel進行多GPU分布式模型訓練
    這個功能實際上早於TensorFlow的第一個版本，早在2012年google內部的前身DistBelief中就已經實現了。這一策略在下圖中得到了很好的說明:在參數伺服器策略中，worker和parameter進程的數量是可變的，每個worker進程在GPU內存中維護自己的模型獨立副本。梯度更新計算如下:在接收到開始信號後，每個工作進程為其特定的批處理片積累梯度。
  • 重磅| Torch7團隊開源PyTorch:Python優先的深度學習框架
    據官網介紹，PyTorch 是一個 Python 優先的深度學習框架，能夠在強大的 GPU 加速基礎上實現張量和動態神經網絡。我們為 GPU 編寫了自定義內存分配器，以保證深度學習模型在運行時有最高的內存效率，這意味著在相同硬體的情況下，你可以訓練比以前更為複雜的深度學習模型。輕鬆拓展編寫新的神經網絡模塊，或與 PyTorch 的 Tensor API 相接的設計都是很直接的，不太抽象。
  • Transformer註解及PyTorch實現(上)
    雖然原文寫得很清楚，但實際上大家普遍反映很難正確地實現。所以我們為此文章寫了篇註解文檔，並給出了一行行實現的Transformer的代碼。本文檔刪除了原文的一些章節並進行了重新排序，並在整個文章中加入了相應的註解。此外，本文檔以Jupyter notebook的形式完成，本身就是直接可以運行的代碼實現，總共有400行庫代碼，在4個GPU上每秒可以處理27,000個tokens。