有bug!用Pytorch Lightning重構代碼速度更慢,修復後速度倍增

2022-01-06 機器之心
用了 Lightning 訓練速度反而更慢，你遇到過這種情況嗎？PyTorch Lightning 是一種重構 PyTorch 代碼的工具，它可以抽出代碼中複雜重複的部分，使得 AI 研究可擴展並且可以快速迭代。然而近日一位名為 Florian Ernst 的博主卻發現 PyTorch Lightning 存在一個 bug——讓原本應該加速的訓練變得更慢了。

Ernst 撰寫博客詳細描述了他發現這個 bug 的過程，以下是博客原文。兩周前，我將一些深度學習代碼重構為 Pytorch Lightning，預計大約有 1.5 倍的加速。然而，訓練、評估和測試任務的速度卻降為原來的 1/4。重構之後的神經網絡需要運行幾天才能得出結果，因此我想找出原因，並儘可能地減少訓練時間。事情是這樣的，我使用的是一些開源深度學習代碼，這些代碼是用來展示某些機器學習任務最新架構的。然而這些代碼本身既不整潔也沒進行優化。我注意到幾個可以加速的地方，並將代碼重構為 Pytorch 代碼，讓訓練大約快了 3 倍。但我認為還有改進的餘地。Pytorch Lightning 是一個非常好的工具：它刪除了大量樣板代碼，並配備了一些優化方法，因此我決定使用 Lightning 重構這些代碼。我原本希望代碼大約能提速 1.5 倍，但完成重構時，我驚訝地發現迭代時間從 4 秒變成了 15 秒，這使訓練時間多了近 3 倍。

我首先運行 Lightning 的分析器來找出問題所在。

基礎分析器給了我一個起點：大部分時間都花在運行一個 epoch 上；高級分析器沒有給我更多信息。我想知道我是否在神經網絡上錯誤地配置了一些超參數。我打亂了其中一些超參數，訓練速度沒有任何變化。然後我調整了數據加載器，發現改變作業數 n_jobs 會對總訓練時間產生影響。然而影響不是加快了計算速度，而是減慢了。

隨著 job 數變化，100 個 epoch 花費的時間。使用 n_jobs=0 完全禁用多處理使我的迭代幾乎比使用 6 個內核快了 2 倍。默認情況下，Pytorch 在兩個 epoch 之間會 kill 掉運行中的進程（worker）並重新加載，因而需要重新加載數據集。在我這個例子中，加載數據集非常慢。我將 DataLoader 裡的 persistent_workers 參數設置為 True，以防止運行中的進程被殺死，進而防止重新加載數據。

# My data Loader parameters
DataLoader(
  train_dataset, batch_size=64, shuffle=True, num_workers=n_workers,
  persistent_workers=True, pin_memory=True,
)

我在 GitHub 上創建了一個 issue，希望 Lightning 團隊意識這個問題，接下來我要尋找問題根源。GitHub 地址：https://github.com/PyTorchLightning/pytorch-lightning/issues/10389Lightning 的 profiler 與上下文管理器一起運行並計算給定塊花費的時間。它可以輕鬆搜索特定的 profiler 操作，以運行「run_training_epoch」為例 。

我開始探究 Lightning 源碼，查看導致循環（loops）變慢的指令，我發現了一些問題：Loop.run 調用 Loop.on_run_start、Loop.on_run_start 重新加載 dataloader，如下圖所示：

Loop.run 調用 Loop.on_run_start…

Loop.on_run_start 重新調用 dataloader問題看起來確實來自在每個 epoch 中重新加載 DataLoader。查看 DataLoader 的源碼，發現是這樣的：

當使用 persistent_workers > 0 迭代 DataLoader 時，如果_iterator` 為 None，則使用_get_iterator() 重新加載整個數據集。可以確定的是 Pytorch Lightning 錯誤地重置了 _iterator，從而導致了這個問題。為了證實這一發現，我用一個自定義的只能重載的__iter__方法替換了 DataLoader：

正如預期的那樣，在迭代之後，_iterator 屬性被正確設置，但在下一個 epoch 開始之前被重置為 None。

n_jobs=1，persistent_workers=True
現在，我只需要知道屬性何時被設置為 None ，這樣就可找到問題的根源。我嘗試使用調試器，但由於多進程或 CUDA 而導致程序崩潰。我開始採用 Python 的 getter & setter 用法：

當 DataLoader._iterator 設置為 None 時，將會列印 stack trace

File "trainer\trainer.py", line 1314, in _run_train
  self.fit_loop.run()
...
File "loops\fit_loop.py", line 234, in advance
  self.epoch_loop.run(data_fetcher)
File "loops\base.py", line 139, in run
  self.on_run_start(*args, **kwargs)
File "loops\epoch\training_epoch_loop.py", line 142, in on_run_start
  self._dataloader_iter = _update_dataloader_iter(...)
File "loops\utilities.py", line 121, in _update_dataloader_iter
  dataloader_iter = enumerate(data_fetcher, batch_idx)
File "utilities\fetching.py", line 198, in __iter__
  self.reset()
File "utilities\fetching.py", line 212, in reset
  self.dataloader.reset()
...
File "trainer\supporters.py", line 498, in _shutdown_workers_and_reset_iterator
  dataloader._iterator = None

通過跟蹤發現每次開始運行時都會調用 DataLoader.reset。通過深入研究代碼後，我發現每次迭代都會重置 DataFetcher，從而導致 DataLoader 也被重置。代碼中沒有條件來避免重置：每個 epoch 都必須重置 DataLoader。既然發現了 bug，就要想辦法修復。修復 bug 非常簡單：我將 self.reset 行從 DataFetcher 的__iter__ 方法中移除：

通過修改後再次訓練，現在一次迭代只需要 1.5 秒，而此前需要 15 秒，使用 vanilla Pytorch 也需要 3 秒，相比較而言，速度確實提升了很多。

我將發現的這個 bug 報告給了 Lightning 團隊，他們對問題進行了修復並在第二天推送了修補程序。我隨後更新了庫，更新後發現他們的修復確實有效。相信更多人將從這次修復中受益，並且他們的 Lightning 模型的訓練和測試時間會得到改善。如果你最近還沒有更新依賴項，請嘗試安裝 pytorch-lightning==1.5.1 或更高版本！原文連結：https://medium.com/@florian-ernst/finding-why-pytorch-lightning-made-my-training-4x-slower-ae64a4720bd1

詳解NVIDIA TAO系列分享第2期：

基於Python的口罩檢測模塊代碼解析——快速搭建基於TensorRT和NVIDIA TAO Toolkit的深度學習訓練環境

第2期線上分享將介紹如何利用NVIDIA TAO Toolkit，在Python的環境下快速訓練並部署一個人臉口罩監測模型，同時會詳細介紹如何利用該工具對模型進行剪枝、評估並優化。TAO Toolkit 內包含了150個預訓練模型，用戶不用從頭開始訓練，極大地減輕了準備樣本的工作量，讓開發者專注於模型的精度提升。本次分享摘要如下：利用TAO Toolkit快速訓練人臉口罩檢測模型

© THE END 

轉載請聯繫本公眾號獲得授權

投稿或尋求報導：content@jiqizhixin.com

相關焦點

  • 【他山之石】pytorch_lightning 全程筆記
    「他山之石，可以攻玉」，站在巨人的肩膀才能看得更高，走得更遠。
  • 分離硬體和代碼、穩定 API,PyTorch Lightning 1.0.0 版本正式發布
    特斯拉 AI 負責人 Andrej Karpathy 也評論稱：「這看起來很棒，也很有前途。PyTorch Lightning 倡導對深度學習代碼進行重構，將『工程（硬體）』與『科學（代碼）』分割開，然後將前者委託給框架。」
  • PyTorch Lightning:專門為機器學習研究者開發的PyTorch輕量 wrapper
    /en/0.7.1/)0.6.0(https://pytorch-lightning.readthedocs.io/en/0.6.0/)0.5.3.2(https://pytorch-lightning.readthedocs.io/en/0.5.3.2/)重構您的PyTorch代碼+好處+完整演練
  • PyTorch Lightning 1.0 正式發布!從0到1,有這9大特點
    Lightning 是為當今世界更複雜的研究和生產案例而設計的，在這些案例中，許多模型使用複雜的規則相互作用。在世界各地，數以百計令人難以置信的工程師和博士們一遍又一遍地實現相同的代碼。Lightning 現在有一個不斷增長的貢獻者社區，其中有超過300個極有才華的深度學習人員，他們選擇分配相同的能量，做完全相同的優化，但是卻有成千上萬的人從他們的努力中受益。
  • 使用 PyTorch Lightning 將深度學習管道速度提高 10 倍
    八年後的 2020 年，微軟 DeepSpeed 團隊在不到 44 分鐘的時間內成功訓練了一個 3.5 億參數的 Large-Bert 模型！九年後，我們現在看到，AlexNet 只是機器學習革命的冰山一角。今天，我們知道許多尚未開發的潛在訓練技術和深度學習模型架構都在我們的掌握之中！
  • 使用PyTorch Lightning自動訓練你的深度神經網絡
    所以，我試圖找到另一個解決方案，然後我找到了PyTorch Lightning，在我看到代碼後，它讓我一見鍾情。因此，我將在本文中介紹的內容是安裝、基本的代碼比較以及通過示例進行比較，這些示例是我自己通過從pytorch lightning site獲取的，一些代碼自己創建的。最後是本文的結論。
  • 【他山之石】Ray和Pytorch Lightning 使用指北
    因為訓練模型，超參搜索，都是沒有什麼創新但是又需要很大的功夫去寫這些重複的代碼。而pytorch lightning這個框架主要解決的就是上述提到的第三步。將模型訓練的大部分代碼整合好，你只需要設定各種參數，即會自動進行訓練，不需要寫各種分布式代碼，各種if else花裡胡哨的。
  • 【他山之石】Pytorch Lightning 完全攻略
    在科研的道路上，更需藉助東風才能更快前行。為此，我們特別搜集整理了一些實用的代碼連結，數據集，軟體，編程技巧等，開闢「他山之石」專欄，助你乘風破浪，一路奮勇向前，敬請關注。第一次發現時，感覺它很重很難學，而且似乎自己也用不上。
  • 從PyTorch到PyTorch Lightning —一個簡要介紹
    但是一旦研究變得複雜，並且將諸如多GPU訓練，16位精度和TPU訓練之類的東西混在一起，用戶很可能會寫出有bug的代碼。PyTorch Lightning完全解決了這個問題。Lightning會構建您的PyTorch代碼，以便抽象出訓練的詳細信息。這使得AI研究可擴展並且可以快速迭代。
  • 實操教程|Pytorch-lightning的使用
    Pytorch-lightning(以下簡稱pl)可以非常簡潔得構建深度學習代碼。但是其實大部分人用不到很多複雜得功能。而pl有時候包裝得過於深了，用的時候稍微有一些不靈活。通常來說，在你的模型搭建好之後，大部分的功能都會被封裝在一個叫trainer的類裡面。
  • PyTorch Lightning工具學習
    PyTorch就是William Falcon在他的博士階段創建的，目標是讓AI研究擴展性更強，忽略一些耗費時間的細節。目前PyTorch Lightning庫已經有了一定的影響力，star已經1w+,同時有超過1千多的研究人員在一起維護這個框架。
  • PyTorch Lightning 工具學習
    簡單介紹PyTorch lightning 是為AI相關的專業的研究人員、研究生、博士等人群開發的。PyTorch就是William Falcon在他的博士階段創建的，目標是讓AI研究擴展性更強，忽略一些耗費時間的細節。
  • 超參搜索調優【Ray】與模型設計工具【Pytorch Lightning 】詳解
    因為訓練模型，超參搜索，都是沒有什麼創新但是又需要很大的功夫去寫這些重複的代碼。而pytorch lightning這個框架主要解決的就是上述提到的第三步。將模型訓練的大部分代碼整合好，你只需要設定各種參數，即會自動進行訓練，不需要寫各種分布式代碼，各種if else花裡胡哨的。
  • 具有PyTorch Lightning的TensorBoard
    因此，數據可視化在使我們的直覺提供更快，更準確的解決方案方面變得極為有用。實際上，數據科學和機器學習日復一日地利用它幾乎所有的機器學習愛好者都可以使用可視化工具。這是我們運行代碼所要滿足的最低要求。使用TensorBoard查看數據Tensorboard文件的保存位置的默認位置是 lightning_logs/ 經過培訓後，在Google Collab筆記本上運行以下命令以打開TensorBoard。
  • 用上Pytorch Lightning的這六招,深度學習pipeline提速10倍!
    而就在最近，一個國外小哥就提出了一種建議：在Pytorch lightning基礎上，讓深度學習pipeline速度提升10倍！用他自己的話來說就是——「爬樓時像給了你一個電梯」。這般「酸爽」，到底是如何做到的呢？
  • 新款MacBook Pro機身過熱時速度變慢 蘋果稱已經修復
    蘋果公司周二證實，該公司新推出的2018年款MacBook Pro筆記本的運行速度會在機身過熱時變慢蘋果公司稱這是個bug，並表示該公司現已推出一個軟體更新來修復這個問題。蘋果公司向財經媒體CNBC發出聲明稱：「在對許多工作負載進行了廣泛的性能測試後，我們已經認定這款固件缺少了一個數位鑰，影響到了溫度管理系統，並可能導致新款MacBook Pro的時鐘頻率在嚴重熱負載的情況下降低。
  • 像 Keras 一樣優雅地使用 pytorch-lightning
    本篇文章為大家介紹一個可以幫助大家優雅地進行深度學習研究的工具：pytorch-lightning。公眾號後臺回復關鍵字：源碼，獲取本文原始碼！pytorch-lightning 是建立在pytorch之上的高層次模型接口，pytorch-lightning之於pytorch,就如同keras之於tensorflow。關於pytorch-lightning的完整入門介紹，可以參考我的另外一篇文章。
  • 【他山之石】pytorch常見的坑匯總
    在科研的道路上，更需藉助東風才能更快前行。為此，我們特別搜集整理了一些實用的代碼連結，數據集，軟體，編程技巧等，開闢「他山之石」專欄，助你乘風破浪，一路奮勇向前，敬請關注。tensorflow有一個比較好用的隊列機制，tf.inputproducer + tfrecord, 但是inputproducer有一個bug，就是無法對每個epoch單獨shuffle，它只能整體shuffle,也就意味著我們無法進行正常的訓練流程（train幾個epoch，在validation上測一個epoch,最終選一個validation上的最好的結果，進行test）。
  • 代碼重構的藝術
    這段代碼可能是別人寫的，也可能時自己寫的，但無論如何，當你覺得這段代碼邏輯糟糕，需要花費幾分鐘才能明白其中的含義時，你就要想著如何去重構才可以使代碼變的更加簡潔直觀有計劃的對代碼重構「找尋重構和開發進度中適合自己的平衡點」 但有些時候我們在準備重構的時會發現，之前的代碼結構和依賴關係錯綜複雜，這時我們就要考慮當前是否有足夠時間去很好的處理這些混亂的代碼
  • PyTorch系列 | 如何加快你的模型訓練速度呢?
    在 GPU 上存儲 Tensors 和運行模型的方法本文的代碼是用 Jupyter notebook，Github 地址為：https://nbviewer.jupyter.org/github/PuneetGrov3r/MediumPosts/blob/master/SpeedUpYourAlgorithms/1%29%20PyTorch.ipynb檢查 cuda 是否可用的代碼非常簡單，