高性能PyTorch是如何煉成的？整理的10條脫坑指南

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

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

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

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

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

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

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

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

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

如果你要處理類圖像的數(shù)據(jù)(2D、3D 掃描)，那么磁盤 I/O 可能會成為瓶頸。為了獲取原始像素?cái)?shù)據(jù)，你的代碼需要從磁盤中讀取數(shù)據(jù)并解碼圖像到內(nèi)存。每個任務(wù)都是迅速的，但是當(dāng)你需要盡快處理成百上千或者成千上萬個任務(wù)時，可能就成了一個挑戰(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 來加載和保存你的訓(xùn)練數(shù)據(jù)，這是從管道中排除最慢的數(shù)據(jù)檢索步驟最簡單的方法。

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

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 實(shí)例，然后將我的訓(xùn)練代碼移到那里。鑒于 V100 比我的 oldie 1080Ti 更新更快的事實(shí)，我期待看到訓(xùn)練快 15–30%。出乎意料的是，每個時期的訓(xùn)練時間都增加了。這讓我明白要注意基礎(chǔ)設(shè)施和環(huán)境差異，而不僅僅是 CPU 和 GPU 的速度。

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

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

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

# 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：離線預(yù)處理所有內(nèi)容

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

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

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

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

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

我們來計(jì)算一下：

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

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

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

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

將 RGB 圖像保持在每個通道深度 8 位?？梢暂p松地在 GPU 上將圖像轉(zhuǎn)換為浮點(diǎn)形式或者標(biāo)準(zhǔ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，減少為原來的五分之一。當(dāng)然，這需要模型中添加額外的步驟來標(biāo)準(zhǔn)化數(shù)據(jù)或?qū)?shù)據(jù)轉(zhuǎn)換為合適的數(shù)據(jù)類型。但是，張量越小，CPU 到 GPU 的傳輸就越快。

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

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

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

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

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

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

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

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

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

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

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

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 上訓(xùn)練圖像分類 pipeline 時，大概可以節(jié)約 20% 的時間。另外值得一提的是，你也可以用 nn.DataParallel 和 nn.DistributedDataParallel 來進(jìn)行推斷。

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

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

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

我編寫的第一個簡單實(shí)現(xiàn)的時候，(除了 bug 之外)花了幾分鐘來計(jì)算單個批的損失值。為了分析 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: 如果設(shè)計(jì)自定義模塊和損失—;—;配置并測試他們

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

時間 VS 金錢

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

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

寫在最后

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

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

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