加入極市專業CV交流群，與?10000+來自港科大、北大、清華、中科院、CMU、騰訊、百度?等名校名企視覺開發者互動交流！

同時提供每月大咖直播分享、真實項目需求對接、干貨資訊匯總，行業技術交流。關注?極市平臺?公眾號?，回復?加群，立刻申請入群~

編譯｜McGL，https://zhuanlan.zhihu.com/p/147723652來源｜https://towardsdatascience.com/efficient-pytorch-part-1-fe40ed5db76c

高效的 PyTorch 訓練pipeline是怎樣的呢？是產生準確率最高模型？還是跑得最快？或是容易理解和擴展？還是很容易并行計算？嗯，以上都是！

在研究和生產領域， PyTorch都是一個很好用的工具。從斯坦福大學、 Udacity、 SalelsForce和Tesla等都采用這個深度學習框架清楚的表明了這一點。然而，每個工具都需要投入時間來最高效地使用和掌握它。在使用 PyTorch 兩年多之后，我決定總結一下我使用這個深度學習庫的經驗。

高效 ——(系統或機器)達到最大生產力，而浪費的精力或費用卻最少。(牛津辭典)

高效 PyTorch 系列的這一部分展示了識別和消除 I/O 和 CPU 、GPU瓶頸的一般技巧。第二部分將揭露一些高效張量運算的技巧。第三部分——高效的模型調試技術。

聲明: 本文假設你至少了解了 PyTorch 的基本知識和概念。

從最明顯的一個開始:

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

nvidia-smi, htop, iotop, nvtop, py-spy, strace 等命令行工具應該成為你最好的朋友。你的訓練pipeline是CPU-bound? IO-bound 還是GPU-bound? 這些工具將幫助你找到答案。

你可能甚至沒有聽說過它們，或者聽說過但沒有使用過。沒關系。如果你不馬上開始使用它們也沒關系。只要記住，別人可能用它們來訓練模型，比你快5%-10%-15%...這最終可能會決定是否贏得或失去市場，得到或失去理想工作崗位的offer。

數據預處理

幾乎所有的訓練pipeline都是從 Dataset 類開始的。它負責提供數據樣本。任何必要的數據轉換和增強都可能在這里發生。簡而言之，Dataset 是一個抽象，它報告數據的大小，并通過給定的索引返回數據樣本。

如果使用類圖像的數據(2D、3D 掃描) ，磁盤I/O可能會成為瓶頸。要獲得原始像素數據，代碼需要從磁盤讀取數據并將圖像解碼到內存中。每個任務都很快，但當你需要盡快處理成千上萬個任務時，這可能會成為一個挑戰。像 NVidia Dali 這樣的庫提供 GPU加速的 JPEG 解碼。如果在數據處理pipeline中遇到 IO 瓶頸，這絕對值得一試。

還有一個選擇。SSD 磁盤的存取時間約為0.08-0.16毫秒。RAM 的訪問時間為納秒。我們可以把我們的數據直接放到內存中！

建議1: 如果可能的話，將所有或部分數據移動到 RAM。

如果你有足夠的內存來加載和保存你所有的訓練數據，這是從pipeline中消除最慢的數據讀取步驟的最簡單的方法。

這個建議對于云實例特別有用，比如 Amazon 的 p3.8 xlarge。此實例有 EBS 磁盤，其默認設置的性能非常有限。然而，這個實例配備了驚人的248Gb 內存。這足夠將所有 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 的家用電腦。有一次我用一個 p3.8xlarge 實例，它有四個 NVidia Tesla V100，我把我的訓練代碼移到了那里。考慮到 V100比我的老款1080Ti 更新更快，我希望看到快15-30% 的訓練速度。令我驚訝的是，每個epoch的訓練時間都在增加！這是我學到的教訓，要注意基礎設施和環境的細微差別，而不僅僅是 CPU 和 GPU 的速度。

根據你的場景，你可以在內存中保持每個文件的二進制內容不變，并“動態”解碼它，或者保留未壓縮的圖像的原始像素。無論你選擇哪種方式，這里有第二個建議:

建議2: 性能分析。測量。比較。每次你對pipeline進行任何改動時，都要仔細評估它對整體的影響。

這個建議僅僅關注訓練速度，假設你不對模型、超參數、數據集等進行更改。你可以擁有一個魔術般的命令行參數(魔術開關) ，如果指定了，它將運行一些合理數量的數據樣本的訓練。有了這個功能，你可以快速的對你的pipeline進行性能分析:

# 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

建議3: 線下預處理所有數據

如果你正在訓練512x512大小的圖像，這些圖像是由2048 × 2048的圖片轉換的，那么事先調整它們的大小。如果你使用灰度圖像作為模型的輸入，請離線進行顏色轉換。如果你正在做 NLP ——事先做tokenization并保存到磁盤。沒有必要在訓練期間一遍又一遍地重復同樣的操作。就漸進式學習而言，你可以保存多種分辨率的訓練數據——這仍然比在線調整目標分辨率要快。

對于表格數據，請考慮在 Dataset 創建時將 pd.DataFrame 對象轉換為 PyTorch 張量。

建議4: 調整 DataLoader 的workers數量

PyTorch 使用 DataLoader 類來簡化為訓練模型生成batches的過程。為了加快速度，它可以并行執行，使用 python 的multiprocessing。大多數情況下，直接用就很好了。以下是一些需要記住的事情:

每個進程生成一批數據，這些batches通過互斥同步(mutex synchronization)提供給主進程。如果你有 N 個workers，那么你的腳本將需要 N 倍的內存才能在系統內存中存儲這些batches。你究竟需要多少RAM？讓我們計算一下:

假設我們用batch size 32 來訓練 Cityscapes 圖像分割模型，RGB 圖像大小為512x512x3 (高, 寬, 通道). 我們在 CPU 端進行圖像標準化(稍后我將解釋為什么它很重要)。在這種情況下，我們的最終圖像張量512 * 512 * 3 * sizeof(float32) = 3,145,728 bytes. 乘以batch size得到的結果是100,663,296 bytes ，即大約100mb

除了圖像，我們需要提供ground-truth masks。它們各自的大小為(默認情況下，掩碼的類型為 long，為8字節)ー 512 * 512 * 1 * 8 * 32 = 67,108,864 , 即大約67mb

因此，一批數據所需的總內存為167 Mb。如果我們有8個workers，所需的總內存量將是167 Mb * 8 = 1,336 Mb.

聽起來還不算太糟，對吧？當你的硬件配置能夠處理超過8個workers所能提供的batches時，問題就出現了。你可以簡單的設置64個workers，但這至少會消耗11gb 的內存。

如果你的數據是3D的，情況會變得更糟; 在這種情況下，即使是一個單通道512x512x512的樣本也將占用134 Mb，而batch size為32將占用4.2 Gb，如果有8個workers，則需要32gb 的內存來保存中間數據。

這個問題有一個部分解決方案ーー你可以盡可能地減少輸入數據的通道深度:

保持 RGB 圖像在每個通道深度為8位。圖像轉換為float和標準化可以很容易地在 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.stdclass 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

這樣做，你可以大大降低內存需求。對于上面的示例，用于內存高效數據表示的內存使用量為每批33 Mb，而不是167 Mb。那是5倍的縮減！當然，這需要在模型本身中執行額外的步驟，將數據normalize/cast為適當的數據類型。然而，張量越小，CPU 到 GPU 的傳輸時間越快。

應該理性的選擇 DataLoader 的workers數量。你應該檢查你的 CPU 和 IO 系統有多快，有多少內存，GPU 處理這些數據有多快。

多GPU訓練與推理

神經網絡模型變得越來越大。今天的趨勢是使用多個 GPU 來提速訓練。由于更大的batch size，它還經常可以改善模型的性能。PyTorch 只需要幾行代碼就可以實現多GPU功能。然而，一些說明乍一看并不明顯。

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

使用多 GPU 最簡單的方法是將模型包在nn.DataParallel類中。在大多數情況下，它工作得不錯，除非你訓練一些圖像分割模型(或任何其它模型，產生大型張量作為輸出)。在前傳結束的時候， nn.DataParallel從所有 GPU 收集輸出回到主 GPU 上，通過輸出后向傳播并進行梯度更新。

有兩個問題:

• GPUs負載不平衡
• 在主 GPU 上收集需要額外的內存

首先，只有主 GPU 在進行損失計算、后向傳播和梯度更新，而其它 GPU 則在60C 涼快處等待下一組數據。

其次，主 GPU 收集所有輸出所需的額外內存通常會迫使你減少batch size。問題是nn.DataParallel 將一批數據均勻的分給各個GPU。假設你有4個GPUs，總batch size為32。那么每個 GPU 將得到8個樣本。但問題是，雖然所有非主 GPU 都可以輕松地將這些batch放入相應的 VRAM 中，但主 GPU 必須分配額外的空間，以保持batch size為32的其它卡的輸出。

GPU使用率不均衡有兩種解決方案:

繼續使用nn.DataParallel并在訓練前傳中計算損失。在這種情況下，你不會將密集預測masks返回給主 GPU，而只返回單個標量損失

使用分布式訓練，也就是nn.DistributedDataParallel. 在分布式訓練的幫助下，你可以解決上面這兩個問題，同時享受觀看所有GPU的100% 負載的樂趣

建議5: 如果你有超過2個 GPU ——考慮使用分布式訓練模式

節省多少時間很大程度上取決于你的場景，根據我的觀察，在4x1080Ti 上訓練圖像分類pipeline時，時間減少了20% 。

另外值得一提的是，你也可以使用nn.DataParallel and nn.DistributedDataParallel來進行推理。

關于自定義損失函數

編寫自定義損失函數是一個有趣和令人興奮的練習。我建議每個人都不時地嘗試一下。實現一個邏輯復雜的損失函數時，有一件事你必須記住: 它運行在 CUDA上，編寫高效的 CUDA 代碼是你的職責。CUDA 高效意味著“沒有 python 控制流”。在 CPU 和 GPU 之間來回切換，訪問 GPU 張量的值可以完成任務，但是性能會很差。

不久前，我實現了一個自定義的cosine embedding損失函數，用于實例分割，該函數來自論文“Segmenting and tracking cell instances with cosine embeddings and recurrent hourglass networks”。它的文本形式相當簡單，但是實現起來有點復雜。

我編寫的第一個簡單的實現(除了 bugs)花了幾分鐘(!) 計算一個batch的損失值。為了分析 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"))

建議6: 如果你設計自定義模塊和損失函數——進行性能分析和測試

在分析了我的初始實現之后，我可以將實現的速度提高100倍。關于在 PyTorch 中編寫高效張量表達式的更多內容將在高效PyTorch 的第二部分中解釋。

時間 vs 金錢

最后但并非最不重要的一點是，有時值得投資功能更強大的硬件，而不是優化代碼。軟件優化總是一個具有不確定結果的高風險過程。升級 CPU，內存，GPU，或所有可能更有效。資金和工程時間都是資源，合理利用兩者是成功的關鍵。

建議7: 有些瓶頸可以通過硬件升級更容易的解決

總結

最大限度的利用你的日常工具是熟練的關鍵。盡量不要走捷徑，如果你不清楚某些事情，就要深入挖掘。總有機會獲得新知識的。問問你自己或你的同事——“我的代碼如何改進? ” 我相信這種追求完美的意識對于計算機工程師來說和其它技能一樣重要。

推薦閱讀

• 給訓練踩踩油門：編寫高效的PyTorch代碼技巧

• Pytorch數據加載的分析

• Pytorch有什么節省內存(顯存)的小技巧？

添加極市小助手微信(ID : cv-mart)，備注：研究方向-姓名-學校/公司-城市(如：目標檢測-小極-北大-深圳)，即可申請加入極市技術交流群，更有每月大咖直播分享、真實項目需求對接、求職內推、算法競賽、干貨資訊匯總、行業技術交流，一起來讓思想之光照的更遠吧~

△長按添加極市小助手

△長按關注極市平臺，獲取最新CV干貨

覺得有用麻煩給個在看啦~??

總結

以上是生活随笔為你收集整理的pytorch dataset读取数据流程_高效 PyTorch ：如何消除训练瓶颈的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。