如何用最少的精力，完成最高效的 PyTorch 訓練?一位有著 PyTorch 兩年使用經歷的 Medium 博主最近分享了他在這方面的 10 個真誠建議。
在 Efficient PyTorch 這一部分中，作者提供了一些識別和消除 I/O 和 CPU 瓶頸的技巧。第二部分闡述了一些高效張量運算的技巧，第三部分是在高效模型上的 debug 技巧。
在閱讀這篇文章之前，你需要對 PyTorch 有一定程度的了解。
好吧，從最明顯的一個開始：
建議 0：了解你代碼中的瓶頸在哪裡
命令行工具比如 nvidia-smi、htop、iotop、nvtop、py-spy、strace 等，應該成為你最好的夥伴。你的訓練管道是否受 CPU 約束?IO 約束?GPU 約束?這些工具將幫你找到答案。
這些工具你可能從未聽過，即使聽過也可能沒用過。沒關係。如果你不立即使用它們也可以。只需記住，其他人可能正在用它們來訓練模型，速度可能會比你快 5%、10%、15%-…… 最終可能會導致面向市場或者工作機會時候的不同結果。
數據預處理
幾乎每個訓練管道都以 Dataset 類開始。它負責提供數據樣本。任何必要的數據轉換和擴充都可能在此進行。簡而言之，Dataset 能報告其規模大小以及在給定索引時，給出數據樣本。
如果你要處理類圖像的數據(2D、3D 掃描)，那麼磁碟 I/O 可能會成為瓶頸。為了獲取原始像素數據，你的代碼需要從磁碟中讀取數據並解碼圖像到內存。每個任務都是迅速的，但是當你需要儘快處理成百上千或者成千上萬個任務時，可能就成了一個挑戰。像 NVidia 這樣的庫會提供一個 GPU 加速的 JPEG 解碼。如果你在數據處理管道中遇到了 IO 瓶頸，這種方法絕對值得一試。
還有另外一個選擇，SSD 磁碟的訪問時間約為 0.08–0.16 毫秒。RAM 的訪問時間是納秒級別的。我們可以直接將數據存入內存。
建議 1：如果可能的話，將數據的全部或部分移至 RAM。
如果你的內存中有足夠多的 RAM 來加載和保存你的訓練數據，這是從管道中排除最慢的數據檢索步驟最簡單的方法。
這個建議可能對雲實例特別有用，比如亞馬遜的 p3.8xlarge。該實例有 EBS 磁碟，它的性能在默認設置下非常受限。但是，該實例配備了驚人的 248Gb 的 RAM。這足夠將整個 ImageNet 數據集存入內存了!你可以通過以下方法達到這一目標：
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%。出乎意料的是，每個時期的訓練時間都增加了。這讓我明白要注意基礎設施和環境差異，而不僅僅是 CPU 和 GPU 的速度。
根據你的方案，你可以將每個文件的二進位內容保持不變，並在 RAM 中進行即時解碼，或者對未壓縮的圖像進行講解碼，並保留原始像素。但是無論你採用什麼方法，這裡有第二條建議：
建議 2：解析、度量、比較。每次你在管道中提出任何改變，要深入地評估它全面的影響。
假設你對模型、超參數和數據集等沒做任何改動，這條建議只關注訓練速度。你可以設置一個魔術命令行參數(魔術開關)，在指定該參數時，訓練會在一些合理的數據樣例上運行。利用這個特點，你可以迅速解析管道。
# 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)，並存入磁碟。在訓練期間一次次重複相同的操作沒有意義。在進行漸進式學習時，你可以以多種解析度保存訓練數據的，這還是比線上調至目標解析度更快。
對於表格數據，請考慮在創建 Dataset 時將 pd.DataFrame 目標轉換為 PyTorch 張量。
建議 4：調整 DataLoader 的工作程序
PyTorch 使用一個 DataLoader 類來簡化用於訓練模型的批處理過程。為了加快速度，它可以使用 Python 中的多進程並行執行。大多數情況下，它可以直接使用。還有幾點需要記住：
每個進程生成一批數據，這些批通過互斥鎖同步可用於主進程。如果你有 N 個工作程序，那麼你的腳本將需要 N 倍的 RAM 才能在系統內存中存儲這些批次的數據。具體需要多少 RAM 呢?
我們來計算一下：
假設我們為 Cityscapes 訓練圖像分割模型，其批處理大小為 32，RGB 圖像大小是 512x512x3(高、寬、通道)。我們在 CPU 端進行圖像標準化(稍後我將會解釋為什麼這一點比較重要)。在這種情況下，我們最終的圖像 tensor 將會是 512 * 512 * 3 * sizeof(float32) = 3,145,728 字節。與批處理大小相乘，結果是 100,663,296 字節，大約 100Mb; 除了圖像之外，我們還需要提供 ground-truth 掩膜。它們各自的大小為(默認情況下，掩膜的類型是 long，8 個字節)——512 * 512 * 1 * 8 * 32 = 67,108,864 或者大約 67Mb; 因此一批數據所需要的總內存是 167Mb。假設有 8 個工作程序，內存的總需求量將是 167 Mb * 8 = 1,336 Mb。
聽起來沒有很糟糕，對嗎?當你的硬體設置能夠容納提供 8 個以上的工作程序提供的更多批處理時，就會出現問題。或許可以天真地放置 64 個工作程序，但是這將消耗至少近 11Gb 的 RAM。
當你的數據是 3D 立體掃描時，情況會更糟糕。在這種情況下，512x512x512 單通道 volume 就會佔 134Mb，批處理大小為 32 時，8 個工作程序將佔 4.2Gb，僅僅是在內存中保存中間數據，你就需要 32Gb 的 RAM。
對於這個問題，有個能解決部分問題的方案——你可以儘可能地減少輸入數據的通道深度：
將 RGB 圖像保持在每個通道深度 8 位。可以輕鬆地在 GPU 上將圖像轉換為浮點形式或者標準化。 在數據集中用 uint8 或 uint16 數據類型代替 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 的需求。對於上面的示例。用於高效存儲數據表示的內存使用量將為每批 33Mb，而之前是 167Mb，減少為原來的五分之一。當然，這需要模型中添加額外的步驟來標準化數據或將數據轉換為合適的數據類型。但是，張量越小，CPU 到 GPU 的傳輸就越快。
DataLoader 的工作程序的數量應該謹慎選擇。你應該查看你的 CPU 和 IO 系統有多快，你有多少內存，GPU 處理數據有多快。
多 GPU 訓練 & 推理神經網絡模型變得越來越大。今天，使用多個 GPU 來增加訓練時間已成為一種趨勢。幸運的是，它經常會提升模型性能來達到更大的批處理量。PyTorch 僅用幾行代碼就可以擁有運行多 GPU 的所有功能。但是，乍一看，有些注意事項並不明顯。
model = nn.DataParallel(model) # Runs model on all available GPUs
運行多 GPU 最簡單的方法就是將模型封裝在 nn.DataParallel 類中。除非你要訓練圖像分割模型(或任何生成大型張量作為輸出的其他模型)，否則大多數情況下效果不錯。在正向推導結束時，nn.DataParallel 將收集主 GPU 上所有的 GPU 輸出，來通過輸出反向運行，並完成梯度更新。
於是，現在就有兩個問題：
GPU 負載不平衡; 在主 GPU 上聚合需要額外的視頻內存
首先，只有主 GPU 能進行損耗計算、反向推導和漸變步驟，其他 GPU 則會在 60 攝氏度以下冷卻，等待下一組數據。
其次，在主 GPU 上聚合所有輸出所需的額外內存通常會促使你減小批處理的大小。nn.DataParallel 將批處理均勻地分配到多個 GPU。假設你有 4 個 GPU，批處理總大小為 32;然後，每個 GPU 將獲得包含 8 個樣本的塊。但問題是，儘管所有的主 GPU 都可以輕鬆地將這些批處理放入對應的 VRAM 中，但主 GPU 必須分配額外的空間來容納 32 個批處理大小，以用於其他卡的輸出。
對於這種不均衡的 GPU 使用率，有兩種解決方案：
在訓練期間繼續在前向推導內使用 nn.DataParallel 計算損耗。在這種情況下。za 不會將密集的預測掩碼返回給主 GPU，而只會返回單個標量損失; 使用分布式訓練，也稱為 nn.DistributedDataParallel。藉助分布式訓練的另一個好處是可以看到 GPU 實現 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
能節省多少時間很大程度上取決於你的方案，我觀察到，在 4x1080Ti 上訓練圖像分類 pipeline 時，大概可以節約 20% 的時間。另外值得一提的是，你也可以用 nn.DataParallel 和 nn.DistributedDataParallel 來進行推斷。
關於自定義損失函數
編寫自定義損失函數是一項很有趣的練習，我建議大家都不時嘗試一下。提到這種邏輯複雜的損失函數，你要牢記一件事：它們都在 CUDA 上運行，你應該會寫「CUDA-efficient」代碼。「CUDA-efficient」意味著「沒有 Python 控制流」。在 CPU 和 GPU 之間來回切換，訪問 GPU 張量的個別值也可以完成這些工作，但是性能表現會很差。
前段時間，我實現了一個自定義餘弦嵌入損失函數，是從《Segmenting and tracking cell instances with cosine embeddings and recurrent hourglass networks》這篇論文中來的，從文本形式上看它非常簡單，但實現起來卻有些複雜。
我編寫的第一個簡單實現的時候，(除了 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: 如果設計自定義模塊和損失——配置並測試他們
在對最初的實現進行性能分析之後，就能夠提速 100 倍。關於在 PyTorch 中編寫高效張量表達式的更多信息，將在 Efficient PyTorch — Part 2 進行說明。
時間 VS 金錢
最後但非常重要的一點，有時候投資功能更強大的硬體，比優化代碼可能更有價值。軟體優化始終是結果無法確定的高風險之旅，升級 CPU、RAM、GPU 或者同時升級以上硬體可能會更有效果。金錢和時間都是資源，二者的均衡利用是成功的關鍵。
通過硬體升級可以更輕鬆地解決某些瓶頸。
寫在最後
懂得充分利用日常工具是提高熟練度的關鍵，儘量不要製造「捷徑」，如果遇到不清楚的地方，請深入挖掘，總有機會發現新知識。正所謂「每日一省」：問問自己，我的代碼還能改進嗎?這種精益求精的信念和其他技能一樣，都是計算機工程師之路的必備品。
原文連結：https://towardsdatascience.com/efficient-pytorch-part-1-fe40ed5db76c
【本文是51CTO專欄機構「機器之心」的原創譯文，微信公眾號「機器之心( id: almosthuman2014)」】