點擊上方“小白學視覺”，選擇加"星標"或“置頂”

重磅干貨，第一時間送達

本文轉(zhuǎn)載自：機器之心 | 作者：Eugene Khvedchenya

高性能 PyTorch 的訓練管道是什么樣的？是產(chǎn)生最高準確率的模型？是最快的運行速度？是易于理解和擴展？還是容易并行化？答案是，包括以上提到的所有。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

建議 3：離線預處理所有內(nèi)容

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

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

建議 4：調(diào)整 DataLoader 的工作程序

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

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

我們來計算一下：

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

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

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

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

將 RGB 圖像保持在每個通道深度 8 位。可以輕松地在 GPU 上將圖像轉(zhuǎn)換為浮點形式或者標準化。
在數(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ù)表示的內(nèi)存使用量將為每批 33Mb，而之前是 167Mb，減少為原來的五分之一。當然，這需要模型中添加額外的步驟來標準化數(shù)據(jù)或?qū)?shù)據(jù)轉(zhuǎn)換為合適的數(shù)據(jù)類型。但是，張量越小，CPU 到 GPU 的傳輸就越快。

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

多 GPU 訓練 & 推理

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

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

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

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

GPU 負載不平衡；
在主 GPU 上聚合需要額外的視頻內(nèi)存

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

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

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

在訓練期間繼續(xù)在前向推導內(nèi)使用 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 提供了一個非常方便的內(nèi)置分析器，非常簡單好用，提供了解決代碼瓶頸的所有信息：

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)化始終是結(jié)果無法確定的高風險之旅，升級 CPU、RAM、GPU 或者同時升級以上硬件可能會更有效果。金錢和時間都是資源，二者的均衡利用是成功的關鍵。

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

寫在最后

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

下載1：OpenCV-Contrib擴展模塊中文版教程
在「小白學視覺」公眾號后臺回復：擴展模塊中文教程，即可下載全網(wǎng)第一份OpenCV擴展模塊教程中文版，涵蓋擴展模塊安裝、SFM算法、立體視覺、目標跟蹤、生物視覺、超分辨率處理等二十多章內(nèi)容。

下載2：Python視覺實戰(zhàn)項目52講
在「小白學視覺」公眾號后臺回復：Python視覺實戰(zhàn)項目，即可下載包括圖像分割、口罩檢測、車道線檢測、車輛計數(shù)、添加眼線、車牌識別、字符識別、情緒檢測、文本內(nèi)容提取、面部識別等31個視覺實戰(zhàn)項目，助力快速學校計算機視覺。

下載3：OpenCV實戰(zhàn)項目20講
在「小白學視覺」公眾號后臺回復：OpenCV實戰(zhàn)項目20講，即可下載含有20個基于OpenCV實現(xiàn)20個實戰(zhàn)項目，實現(xiàn)OpenCV學習進階。

交流群

歡迎加入公眾號讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動駕駛、計算攝影、檢測、分割、識別、醫(yī)學影像、GAN、算法競賽等微信群（以后會逐漸細分），請掃描下面微信號加群，備注：”昵稱+學校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請按照格式備注，否則不予通過。添加成功后會根據(jù)研究方向邀請進入相關微信群。請勿在群內(nèi)發(fā)送廣告，否則會請出群，謝謝理解~

10條PyTorch避坑指南