PyTorch系列 | 如何加快你的模型訓練速度呢?

原題 | Speed Up your Algorithms Part 1 — PyTorch
譯者 | kbsc13("算法猿的成長"公眾號作者)原文 | https://towardsdatascience.com/speed-up-your-algorithms-part-1-pytorch-56d8a4ae7051聲明 | 翻譯是出於交流學習的目的，歡迎轉載，但請保留本文出於，請勿用作商業或者非法用途本文將主要介紹如何採用 cuda 和 pycuda 檢查、初始化 GPU 設備，並讓你的算法跑得更快。
PyTorch 是 torch 的 python 版本，它是 Facebook AI 研究組開發並開源的一個深度學習框架，也是目前非常流行的框架，特別是在研究人員中，短短幾年已經有追上 Tensorflow 的趨勢了。這主要是因為其簡單、動態計算圖的優點。pycuda 是一個 python 第三方庫，用於處理 Nvidia 的 CUDA 並行計算 API 。在 GPU 上存儲 Tensors 和運行模型的方法本文的代碼是用 Jupyter notebook，Github 地址為：https://nbviewer.jupyter.org/github/PuneetGrov3r/MediumPosts/blob/master/SpeedUpYourAlgorithms/1%29%20PyTorch.ipynb
檢查 cuda 是否可用的代碼非常簡單，如下所示：

import torchtorch.cuda.is_available()# True
獲取基本的設備信息，採用 torch.cuda 即可，但如果想得到更詳細的信息，需要採用 pycuda 。
import torchimport pycuda.driver as cudacuda.init()torch.cuda.current_device()cuda.Device(0).name() 
torch.cuda.get_device_name(0) # Get name device with ID '0'# 'Tesla K80'
import pycuda.driver as cudaimport pycuda.autoinit cuda.init() 
class aboutCudaDevices():    def __init__(self):        pass
    def num_devices(self):        """返回 cuda 設備的數量"""        return cuda.Device.count()
    def devices(self):        """獲取所有可用的設備的名稱"""        num = cuda.Device.count()        print("%d device(s) found:"%num)        for i in range(num):            print(cuda.Device(i).name(), "(Id: %d)"%i)
    def mem_info(self):        """獲取所有設備的總內存和可用內存"""        available, total = cuda.mem_get_info()        print("Available: %.2f GB\nTotal:     %.2f GB"%(available/1e9, total/1e9))
    def attributes(self, device_id=0):        """返回指定 id 的設備的屬性信息"""        return cuda.Device(device_id).get_attributes()
    def __repr__(self):        """輸出設備的數量和其id、內存信息"""        num = cuda.Device.count()        string = ""        string += ("%d device(s) found:\n"%num)        for i in range(num):            string += ( "    %d) %s (Id: %d)\n"%((i+1),cuda.Device(i).name(),i))            string += ("          Memory: %.2f GB\n"%(cuda.Device(i).total_memory()/1e9))        return string
aboutCudaDevices()
import torchtorch.cuda.memory_allocated()torch.cuda.memory_cached()
但需要注意的是，上述函數並不會釋放被 tensors 佔用的 GPU 內存，因此並不能增加當前可用的 GPU 內存。
  
在 GPU 上存儲 Tensors 和運行模型的方法
如果是想存儲變量在 cpu 上，可以按下面代碼所示這麼寫：
a = torch.DoubleTensor([1., 2.])
變量 a 將保持在 cpu 上，並在 cpu 上進行各種運算，如果希望將它轉換到 gpu 上，需要採用 .cuda ，可以有以下兩種實現方法
a = torch.FloatTensor([1., 2.]).cuda()a = torch.cuda.FloatTensor([1., 2.])
這種做法會選擇默認的第一個 GPU，查看方式有下面兩種：
torch.cuda.current_device()
a.get_device()
另外，也可以在 GPU 上運行模型，例子如下所示，簡單使用 nn.Sequential 定義一個模型：
sq = nn.Sequential(         nn.Linear(20, 20),         nn.ReLU(),         nn.Linear(20, 4),         nn.Softmax())
怎麼判斷模型是否在 GPU 上運行呢，可以通過下述方法查看模型的參數是否在 GPU 上來判斷：
# From the discussions here: discuss.pytorch.org/t/how-to-check-if-model-is-on-cuda# 參考 https://discuss.pytorch.org/t/how-to-check-if-model-is-on-cuda/180
next(model.parameters()).is_cuda# True
假設有 3 個 GPU ，我們可以初始化和分配 tensors 到任意一個指定的 GPU 上，代碼如下所示，這裡分配 tensors 到指定 GPU 上，有 3 種方法：
cuda0 = torch.device('cuda:0')cuda1 = torch.device('cuda:1')cuda2 = torch.device('cuda:2')
x = torch.Tensor([1., 2.], device=cuda1)x = torch.Tensor([1., 2.]).to(cuda1)x = torch.Tensor([1., 2.]).cuda(cuda1)
torch.cuda.set_device(2) import osos.environ["CUDA_VISIBLE_DEVICES"] = "0,2"
當你有多個 GPU 的時候，就可以將應用的工作劃分，但這裡存在相互之間交流的問題，不過如果不需要頻繁的交換信息，那麼這個問題就可以忽略。實際上，還有另一個問題，在 PyTorch 中所有 GPU 的運算默認都是異步操作。但在 CPU 和 GPU 或者兩個 GPU 之間的數據複製是需要同步的，當你通過函數 torch.cuda.Stream() 創建自己的流時，你必須注意這個同步問題。
cuda = torch.device('cuda')s = torch.cuda.Stream()  A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)with torch.cuda.stream(s):            B = torch.sum(A)
在多模型中，每個 GPU 應用單獨一個模型，並且各自有預處理操作都完成好的一份數據拷貝；每個 GPU 採用切片輸入和模型的拷貝，每個 GPU 將單獨計算結果，並將結果都發送到同一個 GPU 上進行進一步的運算操作。
數據並行的操作要求我們將數據劃分成多份，然後發送給多個 GPU 進行並行的計算。
PyTorch 中實現數據並行的操作可以通過使用 torch.nn.DataParallel。下面是一個簡單的示例。要實現數據並行，第一個方法是採用 nn.parallel 中的幾個函數，分別實現的功能如下所示：複製(Replicate)：將模型拷貝到多個 GPU 上；分發(Scatter)：將輸入數據根據其第一個維度(通常就是 batch 大小)劃分多份，並傳送到多個 GPU 上；收集(Gather)：從多個 GPU 上傳送回來的數據，再次連接回一起；並行的應用(parallel_apply)：將第三步得到的分布式的輸入數據應用到第一步中拷貝的多個模型上。
replicas = nn.parallel.replicate(module, device_ids)inputs = nn.parallel.scatter(input, device_ids)outputs = nn.parallel.parallel_apply(replicas, inputs)result = nn.parallel.gather(outputs, output_device)
實際上，還有一個更簡單的也是常用的實現方法，如下所示，只需一行代碼即可：
model = nn.DataParallel(model, device_ids=device_ids)result = model(input)
根據文章 https://medium.com/@iliakarmanov/multi-gpu-rosetta-stone-d4fa96162986 以及 Github：https://github.com/ilkarman/DeepLearningFrameworks 得到的不同框架在採用單個 GPU 和 4 個 GPU 時運算速度的對比結果，如下所示：
從圖中可以看到數據並行操作儘管存在多 GPU 之間交流的問題，但是提升的速度還是很明顯的。而 PyTorch 的運算速度僅次於 Chainer ，但它的數據並行方式非常簡單，一行代碼即可實現。
torch.multiprocessing 是對 Python 的 multiprocessing 模塊的一個封裝，並且百分比兼容原始模塊，也就是可以採用原始模塊中的如 Queue 、Pipe、Array 等方法。並且為了加快速度，還添加了一個新的方法--share_memory_()，它允許數據處於一種特殊的狀態，可以在不需要拷貝的情況下，任何進程都可以直接使用該數據。
通過該方法，可以共享 Tensors 、模型的參數 parameters ，可以在 CPU 或者 GPU 之間共享它們。
import torch.multiprocessing as mpdef train(model):    for data, labels in data_loader:        optimizer.zero_grad()        loss_fn(model(data), labels).backward()        optimizer.step()  model = nn.Sequential(nn.Linear(n_in, n_h1),                      nn.ReLU(),                      nn.Linear(n_h1, n_out))model.share_memory() processes = []for i in range(4):     p = mp.Process(target=train, args=(model,))    p.start()    processes.append(p)for p in processes:     p.join()
https://pytorch.org/docs/stable/distributed.htmlhttps://documen.tician.de/pycuda/https://pytorch.org/docs/stable/notes/cuda.htmlhttps://discuss.pytorch.org/t/how-to-check-if-model-is-on-cudahttps://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.htmlhttps://medium.com/@iliakarmanov/multi-gpu-rosetta-stone-d4fa96162986