如何用最少的精力，完成最高效的 PyTorch 訓練?一位有著 PyTorch 兩年使用經(jīng)歷的 Medium 博主最近分享了他在這方面的 10 個真誠建議。

在 Efficient PyTorch 這一部分中，作者提供了一些識別和消除 I/O 和 CPU 瓶頸的技巧。第二部分闡述了一些高效張量運算的技巧，第三部分是在高效模型上的 debug 技巧。

在閱讀這篇文章之前，你需要對 PyTorch 有一定程度的了解。

[[330988]]

好吧，從最明顯的一個開始：

建議 0：了解你代碼中的瓶頸在哪里

命令行工具比如 nvidia-smi、htop、iotop、nvtop、py-spy、strace 等，應該成為你最好的伙伴。你的訓練管道是否受 CPU 約束?IO 約束?GPU 約束?這些工具將幫你找到答案。

這些工具你可能從未聽過，即使聽過也可能沒用過。沒關系。如果你不立即使用它們也可以。只需記住，其他人可能正在用它們來訓練模型，速度可能會比你快 5%、10%、15%-…… 最終可能會導致面向市場或者工作機會時候的不同結果。

數(shù)據(jù)預處理

幾乎每個訓練管道都以 Dataset 類開始。它負責提供數(shù)據(jù)樣本。任何必要的數(shù)據(jù)轉換和擴充都可能在此進行。簡而言之，Dataset 能報告其規(guī)模大小以及在給定索引時，給出數(shù)據(jù)樣本。

如果你要處理類圖像的數(shù)據(jù)(2D、3D 掃描)，那么磁盤 I/O 可能會成為瓶頸。為了獲取原始像素數(shù)據(jù)，你的代碼需要從磁盤中讀取數(shù)據(jù)并解碼圖像到內存。每個任務都是迅速的，但是當你需要盡快處理成百上千或者成千上萬個任務時，可能就成了一個挑戰(zhàn)。像 NVidia 這樣的庫會提供一個 GPU 加速的 JPEG 解碼。如果你在數(shù)據(jù)處理管道中遇到了 IO 瓶頸，這種方法絕對值得一試。

還有另外一個選擇，SSD 磁盤的訪問時間約為 0.08–0.16 毫秒。RAM 的訪問時間是納秒級別的。我們可以直接將數(shù)據(jù)存入內存。

建議 1：如果可能的話，將數(shù)據(jù)的全部或部分移至 RAM。

如果你的內存中有足夠多的 RAM 來加載和保存你的訓練數(shù)據(jù)，這是從管道中排除最慢的數(shù)據(jù)檢索步驟最簡單的方法。

這個建議可能對云實例特別有用，比如亞馬遜的 p3.8xlarge。該實例有 EBS 磁盤，它的性能在默認設置下非常受限。但是，該實例配備了驚人的 248Gb 的 RAM。這足夠將整個 ImageNet 數(shù)據(jù)集存入內存了!你可以通過以下方法達到這一目標：

class RAMDataset(Dataset): 
  def __init__(image_fnames, targets): 
    self.targets = targets 
    self.images = [] 
    for fname in tqdm(image_fnames, desc="Loading files in RAM"): 
      with open(fname, "rb") as f: 
        self.images.append(f.read()) 
 
  def __len__(self): 
    return len(self.targets) 
 
  def __getitem__(self, index): 
    target = self.targets[index] 
    image, retval = cv2.imdecode(self.images[index], cv2.IMREAD_COLOR) 
    return image, target 

我個人也面對過這個瓶頸問題。我有一臺配有 4x1080Ti GPUs 的家用 PC。有一次，我采用了有 4 個 NVidia Tesla V100 的 p3.8xlarge 實例，然后將我的訓練代碼移到那里。鑒于 V100 比我的 oldie 1080Ti 更新更快的事實，我期待看到訓練快 15–30%。出乎意料的是，每個時期的訓練時間都增加了。這讓我明白要注意基礎設施和環(huán)境差異，而不僅僅是 CPU 和 GPU 的速度。

根據(jù)你的方案，你可以將每個文件的二進制內容保持不變，并在 RAM 中進行即時解碼，或者對未壓縮的圖像進行講解碼，并保留原始像素。但是無論你采用什么方法，這里有第二條建議：

建議 2：解析、度量、比較。每次你在管道中提出任何改變，要深入地評估它全面的影響。

假設你對模型、超參數(shù)和數(shù)據(jù)集等沒做任何改動，這條建議只關注訓練速度。你可以設置一個魔術命令行參數(shù)(魔術開關)，在指定該參數(shù)時，訓練會在一些合理的數(shù)據(jù)樣例上運行。利用這個特點，你可以迅速解析管道。

# Profile CPU bottlenecks 
python -m cProfile training_script.py --profiling 
 
# Profile GPU bottlenecks 
nvprof --print-gpu-trace python train_mnist.py 
 
# Profile system calls bottlenecks 
strace -fcT python training_script.py -e trace=open,close,read 
 
Advice 3: *Preprocess everything offline* 

建議 3：離線預處理所有內容

如果你要訓練由多張 2048x2048 圖像制成的 512x512 尺寸圖像，請事先調整。如果你使用灰度圖像作為模型的輸入，請離線調整顏色。如果你正在進行自然語言處理(NLP)，請事先做分詞處理(tokenization)，并存入磁盤。在訓練期間一次次重復相同的操作沒有意義。在進行漸進式學習時，你可以以多種分辨率保存訓練數(shù)據(jù)的，這還是比線上調至目標分辨率更快。

對于表格數(shù)據(jù)，請考慮在創(chuàng)建 Dataset 時將 pd.DataFrame 目標轉換為 PyTorch 張量。

建議 4：調整 DataLoader 的工作程序

PyTorch 使用一個 DataLoader 類來簡化用于訓練模型的批處理過程。為了加快速度，它可以使用 Python 中的多進程并行執(zhí)行。大多數(shù)情況下，它可以直接使用。還有幾點需要記住：

每個進程生成一批數(shù)據(jù)，這些批通過互斥鎖同步可用于主進程。如果你有 N 個工作程序，那么你的腳本將需要 N 倍的 RAM 才能在系統(tǒng)內存中存儲這些批次的數(shù)據(jù)。具體需要多少 RAM 呢?

我們來計算一下：

• 假設我們?yōu)?Cityscapes 訓練圖像分割模型，其批處理大小為 32，RGB 圖像大小是 512x512x3(高、寬、通道)。我們在 CPU 端進行圖像標準化(稍后我將會解釋為什么這一點比較重要)。在這種情況下，我們最終的圖像 tensor 將會是 512 * 512 * 3 * sizeof(float32) = 3,145,728 字節(jié)。與批處理大小相乘，結果是 100,663,296 字節(jié)，大約 100Mb;
• 除了圖像之外，我們還需要提供 ground-truth 掩膜。它們各自的大小為(默認情況下，掩膜的類型是 long，8 個字節(jié))——512 * 512 * 1 * 8 * 32 = 67,108,864 或者大約 67Mb;
• 因此一批數(shù)據(jù)所需要的總內存是 167Mb。假設有 8 個工作程序，內存的總需求量將是 167 Mb * 8 = 1,336 Mb。

聽起來沒有很糟糕，對嗎?當你的硬件設置能夠容納提供 8 個以上的工作程序提供的更多批處理時，就會出現(xiàn)問題。或許可以天真地放置 64 個工作程序，但是這將消耗至少近 11Gb 的 RAM。

當你的數(shù)據(jù)是 3D 立體掃描時，情況會更糟糕。在這種情況下，512x512x512 單通道 volume 就會占 134Mb，批處理大小為 32 時，8 個工作程序將占 4.2Gb，僅僅是在內存中保存中間數(shù)據(jù)，你就需要 32Gb 的 RAM。

對于這個問題，有個能解決部分問題的方案——你可以盡可能地減少輸入數(shù)據(jù)的通道深度：

• 將 RGB 圖像保持在每個通道深度 8 位。可以輕松地在 GPU 上將圖像轉換為浮點形式或者標準化。
• 在數(shù)據(jù)集中用 uint8 或 uint16 數(shù)據(jù)類型代替 long。

class MySegmentationDataset(Dataset): 
  ... 
  def __getitem__(self, index): 
    image = cv2.imread(self.images[index]) 
    target = cv2.imread(self.masks[index]) 
 
    # No data normalization and type casting here 
    return torch.from_numpy(image).permute(2,0,1).contiguous(), 
           torch.from_numpy(target).permute(2,0,1).contiguous() 
 
class Normalize(nn.Module): 
    # https://github.com/BloodAxe/pytorch-toolbelt/blob/develop/pytorch_toolbelt/modules/normalize.py 
    def __init__(self, mean, std): 
        super().__init__() 
        self.register_buffer("mean", torch.tensor(mean).float().reshape(1, len(mean), 1, 1).contiguous()) 
        self.register_buffer("std", torch.tensor(std).float().reshape(1, len(std), 1, 1).reciprocal().contiguous()) 
 
    def forward(self, input: torch.Tensor) -> torch.Tensor: 
        return (input.to(self.mean.type) - self.mean) * self.std 
 
class MySegmentationModel(nn.Module): 
  def __init__(self): 
    self.normalize = Normalize([0.221 * 255], [0.242 * 255]) 
    self.loss = nn.CrossEntropyLoss() 
 
  def forward(self, image, target): 
    image = self.normalize(image) 
    output = self.backbone(image) 
 
    if target is not None: 
      loss = self.loss(output, target.long()) 
      return loss 
 
    return output 

通過這樣做，會大大減少 RAM 的需求。對于上面的示例。用于高效存儲數(shù)據(jù)表示的內存使用量將為每批 33Mb，而之前是 167Mb，減少為原來的五分之一。當然，這需要模型中添加額外的步驟來標準化數(shù)據(jù)或將數(shù)據(jù)轉換為合適的數(shù)據(jù)類型。但是，張量越小，CPU 到 GPU 的傳輸就越快。

DataLoader 的工作程序的數(shù)量應該謹慎選擇。你應該查看你的 CPU 和 IO 系統(tǒng)有多快，你有多少內存，GPU 處理數(shù)據(jù)有多快。

多 GPU 訓練 & 推理

[[330989]]

神經(jīng)網(wǎng)絡模型變得越來越大。今天，使用多個 GPU 來增加訓練時間已成為一種趨勢。幸運的是，它經(jīng)常會提升模型性能來達到更大的批處理量。PyTorch 僅用幾行代碼就可以擁有運行多 GPU 的所有功能。但是，乍一看，有些注意事項并不明顯。

model = nn.DataParallel(model) # Runs model on all available GPUs 

運行多 GPU 最簡單的方法就是將模型封裝在 nn.DataParallel 類中。除非你要訓練圖像分割模型(或任何生成大型張量作為輸出的其他模型)，否則大多數(shù)情況下效果不錯。在正向推導結束時，nn.DataParallel 將收集主 GPU 上所有的 GPU 輸出，來通過輸出反向運行，并完成梯度更新。

于是，現(xiàn)在就有兩個問題：

• GPU 負載不平衡;
• 在主 GPU 上聚合需要額外的視頻內存

首先，只有主 GPU 能進行損耗計算、反向推導和漸變步驟，其他 GPU 則會在 60 攝氏度以下冷卻，等待下一組數(shù)據(jù)。

其次，在主 GPU 上聚合所有輸出所需的額外內存通常會促使你減小批處理的大小。nn.DataParallel 將批處理均勻地分配到多個 GPU。假設你有 4 個 GPU，批處理總大小為 32;然后，每個 GPU 將獲得包含 8 個樣本的塊。但問題是，盡管所有的主 GPU 都可以輕松地將這些批處理放入對應的 VRAM 中，但主 GPU 必須分配額外的空間來容納 32 個批處理大小，以用于其他卡的輸出。

對于這種不均衡的 GPU 使用率，有兩種解決方案：

• 在訓練期間繼續(xù)在前向推導內使用 nn.DataParallel 計算損耗。在這種情況下。za 不會將密集的預測掩碼返回給主 GPU，而只會返回單個標量損失;
• 使用分布式訓練，也稱為 nn.DistributedDataParallel。借助分布式訓練的另一個好處是可以看到 GPU 實現(xiàn) 100% 負載。

如果想了解更多，可以看看這三篇文章：

• https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
• https://medium.com/@theaccelerators/learn-pytorch-multi-gpu-properly-3eb976c030ee
• https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2

建議 5: 如果你擁有兩個及以上的 GPU

能節(jié)省多少時間很大程度上取決于你的方案，我觀察到，在 4x1080Ti 上訓練圖像分類 pipeline 時，大概可以節(jié)約 20% 的時間。另外值得一提的是，你也可以用 nn.DataParallel 和 nn.DistributedDataParallel 來進行推斷。

關于自定義損失函數(shù)

編寫自定義損失函數(shù)是一項很有趣的練習，我建議大家都不時嘗試一下。提到這種邏輯復雜的損失函數(shù)，你要牢記一件事：它們都在 CUDA 上運行，你應該會寫「CUDA-efficient」代碼。「CUDA-efficient」意味著「沒有 Python 控制流」。在 CPU 和 GPU 之間來回切換，訪問 GPU 張量的個別值也可以完成這些工作，但是性能表現(xiàn)會很差。

前段時間，我實現(xiàn)了一個自定義余弦嵌入損失函數(shù)，是從《Segmenting and tracking cell instances with cosine embeddings and recurrent hourglass networks》這篇論文中來的，從文本形式上看它非常簡單，但實現(xiàn)起來卻有些復雜。

我編寫的第一個簡單實現(xiàn)的時候，(除了 bug 之外)花了幾分鐘來計算單個批的損失值。為了分析 CUDA 瓶頸，PyTorch 提供了一個非常方便的內置分析器，非常簡單好用，提供了解決代碼瓶頸的所有信息：

def test_loss_profiling(): 
    loss = nn.BCEWithLogitsLoss() 
    with torch.autograd.profiler.profile(use_cuda=True) as prof: 
        input = torch.randn((8, 1, 128, 128)).cuda() 
        input.requires_grad = True 
 
        target = torch.randint(1, (8, 1, 128, 128)).cuda().float() 
 
        for i in range(10): 
            l = loss(input, target) 
            l.backward() 
    print(prof.key_averages().table(sort_by="self_cpu_time_total")) 

建議 9: 如果設計自定義模塊和損失——配置并測試他們

在對最初的實現(xiàn)進行性能分析之后，就能夠提速 100 倍。關于在 PyTorch 中編寫高效張量表達式的更多信息，將在 Efficient PyTorch — Part 2 進行說明。

時間 VS 金錢

最后但非常重要的一點，有時候投資功能更強大的硬件，比優(yōu)化代碼可能更有價值。軟件優(yōu)化始終是結果無法確定的高風險之旅，升級 CPU、RAM、GPU 或者同時升級以上硬件可能會更有效果。金錢和時間都是資源，二者的均衡利用是成功的關鍵。

通過硬件升級可以更輕松地解決某些瓶頸。

寫在最后

懂得充分利用日常工具是提高熟練度的關鍵，盡量不要制造「捷徑」，如果遇到不清楚的地方，請深入挖掘，總有機會發(fā)現(xiàn)新知識。正所謂「每日一省」：問問自己，我的代碼還能改進嗎?這種精益求精的信念和其他技能一樣，都是計算機工程師之路的必備品。

原文鏈接：https://towardsdatascience.com/efficient-pytorch-part-1-fe40ed5db76c

【本文是51CTO專欄機構“機器之心”的原創(chuàng)譯文，微信公眾號“機器之心( id: almosthuman2014)”】 

戳這里，看該作者更多好文