如何用最少的精力，完成高效的 PyTorch 訓(xùn)練?一位有著 PyTorch 兩年使用經(jīng)歷的 Medium 博主最近分享了他在這方面的 10 個真誠建議。
在 Efficient PyTorch 這一部分中，作者提供了一些識別和消除 I/O 和 CPU 瓶頸的技巧。第二部分闡述了一些高效張量運算的技巧，第三部分是在高效模型上的 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)致面向市場或者工作機會時候的不同結(jié)果。
數(shù)據(jù)預(yù)處理
幾乎每個訓(xùn)練管道都以 Dataset 類開始。它負(fù)責(zé)提供數(shù)據(jù)樣本。任何必要的數(shù)據(jù)轉(zhuǎn)換和擴充都可能在此進行。簡而言之，Dataset 能報告其規(guī)模大小以及在給定索引時，給出數(shù)據(jù)樣本。
如果你要處理類圖像的數(shù)據(jù)(2D、3D 掃描)，那么磁盤 I/O 可能會成為瓶頸。為了獲取原始像素數(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ù)檢索步驟最簡單的方法。
這個建議可能對云實例特別有用，比如亞馬遜的 p3.8xlarge。該實例有 EBS 磁盤，它的性能在默認(rèn)設(shè)置下非常受限。但是，該實例配備了驚人的 248Gb 的 RAM。這足夠?qū)⒄麄€ ImageNet 數(shù)據(jù)集存入內(nèi)存了!你可以通過以下方法達(dá)到這一目標(biāo)：
classRAMDataset(Dataset):def__init__(image_fnames,targets):self.targets=targetsself.images=[]forfnameintqdm(image_fnames,desc="LoadingfilesinRAM"):withopen(fname,"rb")asf:self.images.append(f.read())def__len__(self):returnlen(self.targets)def__getitem__(self,index):target=self.targets[index]image,retval=cv2.imdecode(self.images[index],cv2.IMREAD_COLOR)returnimage,target
我個人也面對過這個瓶頸問題。我有一臺配有 4x1080Ti GPUs 的家用 PC。有一次，我采用了有 4 個 NVidia Tesla V100 的 p3.8xlarge 實例，然后將我的訓(xùn)練代碼移到那里。鑒于 V100 比我的 oldie 1080Ti 更新更快的事實，我期待看到訓(xùn)練快 15–30%。出乎意料的是，每個時期的訓(xùn)練時間都增加了。這讓我明白要注意基礎(chǔ)設(shè)施和環(huán)境差異，而不僅僅是 CPU 和 GPU 的速度。
根據(jù)你的方案，你可以將每個文件的二進制內(nèi)容保持不變，并在 RAM 中進行即時解碼，或者對未壓縮的圖像進行講解碼，并保留原始像素。但是無論你采用什么方法，這里有第二條建議：
建議 2：解析、度量、比較。每次你在管道中提出任何改變，要深入地評估它全面的影響。
假設(shè)你對模型、超參數(shù)和數(shù)據(jù)集等沒做任何改動，這條建議只關(guān)注訓(xùn)練速度。你可以設(shè)置一個魔術(shù)命令行參數(shù)(魔術(shù)開關(guān))，在指定該參數(shù)時，訓(xùn)練會在一些合理的數(shù)據(jù)樣例上運行。利用這個特點，你可以迅速解析管道。
#ProfileCPUbottleneckspython-mcProfiletraining_script.py--profiling#ProfileGPUbottlenecksnvprof--print-gpu-tracepythontrain_mnist.py#Profilesystemcallsbottlenecksstrace-fcTpythontraining_script.py-etrace=open,close,readAdvice3:*Preprocesseverythingoffline*
建議 3：離線預(yù)處理所有內(nèi)容
如果你要訓(xùn)練由多張 2048x2048 圖像制成的 512x512 尺寸圖像，請事先調(diào)整。如果你使用灰度圖像作為模型的輸入，請離線調(diào)整顏色。如果你正在進行自然語言處理(NLP)，請事先做分詞處理(tokenization)，并存入磁盤。在訓(xùn)練期間一次次重復(fù)相同的操作沒有意義。在進行漸進式學(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 中的多進程并行執(zhí)行。大多數(shù)情況下，它可以直接使用。還有幾點需要記?。?/p>
每個進程生成一批數(shù)據(jù)，這些批通過互斥鎖同步可用于主進程。如果你有 N 個工作程序，那么你的腳本將需要 N 倍的 RAM 才能在系統(tǒng)內(nèi)存中存儲這些批次的數(shù)據(jù)。具體需要多少 RAM 呢?
我們來計算一下：
假設(shè)我們?yōu)?Cityscapes 訓(xùn)練圖像分割模型，其批處理大小為 32，RGB 圖像大小是 512x512x3(高、寬、通道)。我們在 CPU 端進行圖像標(biāo)準(zhǔ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)換為浮點形式或者標(biāo)準(zhǔn)化。 在數(shù)據(jù)集中用 uint8 或 uint16 數(shù)據(jù)類型代替 long。 classMySegmentationDataset(Dataset):...def__getitem__(self,index):image=cv2.imread(self.images[index])target=cv2.imread(self.masks[index])#Nodatanormalizationandtypecastingherereturntorch.from_numpy(image).permute(2,0,1).contiguous(),torch.from_numpy(target).permute(2,0,1).contiguous()classNormalize(nn.Module):#https://github.com/BloodAxe/pytorch-toolbelt/blob/develop/pytorch_toolbelt/modules/normalize.pydef__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())defforward(self,input:torch.Tensor)->torch.Tensor:return(input.to(self.mean.type)-self.mean)*self.stdclassMySegmentationModel(nn.Module):def__init__(self):self.normalize=Normalize([0.221*255],[0.242*255])self.loss=nn.CrossEntropyLoss()defforward(self,image,target):image=self.normalize(image)output=self.backbone(image)iftargetisnotNone:loss=self.loss(output,target.long())returnlossreturnoutput
通過這樣做，會大大減少 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)練時間已成為一種趨勢。幸運的是，它經(jīng)常會提升模型性能來達(dá)到更大的批處理量。PyTorch 僅用幾行代碼就可以擁有運行多 GPU 的所有功能。但是，乍一看，有些注意事項并不明顯。
model=nn.DataParallel(model)#RunsmodelonallavailableGPUs
運行多 GPU 最簡單的方法就是將模型封裝在 nn.DataParallel 類中。除非你要訓(xùn)練圖像分割模型(或任何生成大型張量作為輸出的其他模型)，否則大多數(shù)情況下效果不錯。在正向推導(dǎo)結(jié)束時，nn.DataParallel 將收集主 GPU 上所有的 GPU 輸出，來通過輸出反向運行，并完成梯度更新。
于是，現(xiàn)在就有兩個問題：
GPU 負(fù)載不平衡; 在主 GPU 上聚合需要額外的視頻內(nèi)存
首先，只有主 GPU 能進行損耗計算、反向推導(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 計算損耗。在這種情況下。za 不會將密集的預(yù)測掩碼返回給主 GPU，而只會返回單個標(biāo)量損失; 使用分布式訓(xùn)練，也稱為 nn.DistributedDataParallel。借助分布式訓(xùn)練的另一個好處是可以看到 GPU 實現(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 來進行推斷。
關(guān)于自定義損失函數(shù)
編寫自定義損失函數(shù)是一項很有趣的練習(xí)，我建議大家都不時嘗試一下。提到這種邏輯復(fù)雜的損失函數(shù)，你要牢記一件事：它們都在 CUDA 上運行，你應(yīng)該會寫「CUDA-efficient」代碼?！窩UDA-efficient」意味著「沒有 Python 控制流」。在 CPU 和 GPU 之間來回切換，訪問 GPU 張量的個別值也可以完成這些工作，但是性能表現(xiàn)會很差。
前段時間，我實現(xiàn)了一個自定義余弦嵌入損失函數(shù)，是從《Segmenting and tracking cell instances with cosine embeddings and recurrent hourglass networks》這篇論文中來的，從文本形式上看它非常簡單，但實現(xiàn)起來卻有些復(fù)雜。
我編寫的第一個簡單實現(xiàn)的時候，(除了 bug 之外)花了幾分鐘來計算單個批的損失值。為了分析 CUDA 瓶頸，PyTorch 提供了一個非常方便的內(nèi)置分析器，非常簡單好用，提供了解決代碼瓶頸的所有信息：
deftest_loss_profiling():loss=nn.BCEWithLogitsLoss()withtorch.autograd.profiler.profile(use_cuda=True)asprof:input=torch.randn((8,1,128,128)).cuda()input.requires_grad=Truetarget=torch.randint(1,(8,1,128,128)).cuda().float()foriinrange(10):l=loss(input,target)l.backward()print(prof.key_averages().table(sort_by="self_cpu_time_total"))
建議 9: 如果設(shè)計自定義模塊和損失——配置并測試他們
在對最初的實現(xiàn)進行性能分析之后，就能夠提速 100 倍。關(guān)于在 PyTorch 中編寫高效張量表達(dá)式的更多信息，將在 Efficient PyTorch — Part 2 進行說明。
時間 VS 金錢
最后但非常重要的一點，有時候投資功能更強大的硬件，比優(yōu)化代碼可能更有價值。軟件優(yōu)化始終是結(jié)果無法確定的高風(fēng)險之旅，升級 CPU、RAM、GPU 或者同時升級以上硬件可能會更有效果。金錢和時間都是資源，二者的均衡利用是成功的關(guān)鍵。
通過硬件升級可以更輕松地解決某些瓶頸。
寫在最后
懂得充分利用日常工具是提高熟練度的關(guān)鍵，盡量不要制造「捷徑」，如果遇到不清楚的地方，請深入挖掘，總有機會發(fā)現(xiàn)新知識。正所謂「每日一省」：問問自己，我的代碼還能改進嗎?這種精益求精的信念和其他技能一樣，都是計算機工程師之路的必備品。
重慶分公司
重慶分公司
