?PaperWeekly 原創 ·?作者｜桑運鑫?

學校｜上海交通大學碩士生?

研究方向｜圖神經網絡應用

最近對一個大規模的圖訓練嵌入，發現相關的中文資料還是很欠缺的，把自己踩的一些坑記下來。本文主要針對 DGL [1]?和 PyTorch?[2]?兩個框架。

訓練大規模圖

對于大規模圖不能像小圖一樣把整張圖扔進去訓練，需要對大圖進行采樣，即通過Neighborhood Sampling方法每次采樣一部分輸出節點，然后把更新它們所需的所有節點作為輸入節點，通過這樣的方式做 mini-batch 迭代訓練。具體的用法可以參考官方文檔中的 Chapter 6: Stochastic Training on Large Graphs [3]。

但是 GATNE-T [4] 中有一種更有趣的做法，即只把 DGL?作為一個輔助計算流的工具，提供Neighborhood Sampling和Message Passing等過程，把Node Embedding和Edge Embedding等存儲在圖之外，做一個單獨的Embedding矩陣。每次從 dgl 中獲取節點的id之后再去Embedding矩陣中去取對應的 embedding 進行優化，以此可以更方便的做一些優化。

縮小圖規模

從圖的Message Passing過程可以看出，基本上所有的圖神經網絡的計算都只能傳播連通圖的信息，所以可以先用?connected_componets [5] 檢查一下自己的圖是否是連通圖。如果分為多個連通子圖的話，可以分別進行訓練或者選擇一個大小適中的 component 訓練。

如果圖還是很大的話，也可以對圖整體做一次Neighborhood Sampling采樣一個子圖進行訓練。

減小內存占用

對于大規模數據而言，如何在內存中存下它也是一件讓人傷腦筋的事情。這時候采用什么樣的數據結構存儲就很關鍵了。首先是不要使用原生的list，使用np.ndarray或者torch.tensor。尤其注意不要顯式的使用set存儲大規模數據（可以使用set去重，但不要存儲它）。

注意：四種數據結構消耗的內存之間的差別（比例關系）會隨著數據規模變大而變大。

其次就是在PyTorch中，設置DataLoader的num_workers大于 0 時會出現內存泄露的問題，內存占用會隨著 epoch 不斷增大。查閱資料有兩個解決方法：

• 根據Num_workers in DataLoader will increase memory usage? [6]，設置num_workers小于實際 cpu 數，親測無效；

• 根據 CPU memory gradually leaks when num_workers > 0 in the DataLoader [7]，將原始list轉為np.ndarray或者torch.tensor，可以解決。原因是：There is no way of storing arbitrary python objects (even simple lists) in shared memory in Python without triggering copy-on-write behaviour due to the addition of refcounts, everytime something reads from these objects.

減小顯存消耗

對于大規模圖嵌入而言，Embedding矩陣會非常大。在反向傳播中如果對整個矩陣做優化的話很可能會爆顯存。可以參考 pinsage [8]?的代碼，設置Embedding矩陣的sparse = True，使用?SparseAdam?[9] 進行優化。SparseAdam是一種為 sparse 向量設計的優化方法，它只優化矩陣中參與計算的元素，可以大大減少backward過程中的顯存消耗。

此外，如果顯存仍然不夠的話，可以考慮將Embedding矩陣放到 CPU 上。使用兩個優化器分別進行優化。

加快訓練速度

對于大規模數據而言，訓練同樣要花很長時間。加快訓練速度也很關鍵。加快訓練速度方面主要在兩個方面：加快數據預處理和提高 GPU 利用率。

在加快數據預處理中，大部分數據集樣本之間都是獨立的，可以并行處理，所以當數據規模很大的時候，一定要加大數據預處理的并行度。

不要使用 for 循環逐條處理，可以使用 multiprocess [10] 庫開多進程并行處理。但是要注意適當設置processes，否則會出現錯誤OSError: [Errno 24] Too many open files。

此外，數據處理好之后最好保存為pickle格式的文件，下次使用可以直接加載，不要再花時間處理一遍。

在提高 GPU 利用率上，如果 GPU 利用率比較低主要是兩個原因：batch_size較小（表現為 GPU 利用率一直很低）和數據加載跟不上訓練（表現為 GPU 利用率上升下降上升下降，如此循環）。解決方法也是兩個：

• 增大batch_size，一般來說 GPU 利用率和batch_size成正比；

• 加快數據加載：設置DataLoader的pin_memory=True, 適當增大num_workers（注意不要盲目增大，設置到使用的 CPU 利用率到 90% 左右就可以了，不然反而可能會因為開線程的消耗拖慢訓練速度）。

如果有多塊 GPU 的話，可以參照 graphsage [11] 的代碼進行多 GPU 訓練。一般來說在單機多卡的情況下都可以得到線性加速。但是使用DistributedDataParallel有幾個需要注意的問題：

• 最好參照 graphsage [11] 中的代碼而不是使用官方教程中的torch.multiprocessing.spawn函數開辟多進程。因為使用這個函數會不停的打印Using backend: pytorch，暫時還不清楚是什么原因。

• 和DataParallel一樣，DistributedDataParallel會對原模型進行封裝，如果需要獲取原模型model的一些屬性或函數的話，需要將model替換為model.module。

• 在使用DistributedDataParallel時，需要根據 GPU 的數量對batch_size和learning rate進行調整。根據 Should we split batch_size according to ngpu_per_node when DistributedDataparallel [12] ，簡單來說就是保持batch_size和learning rate的乘積不變，因為我們多 GPU 訓練一般不改batch_size，所以使用了多少 GPU 就要把learning rate擴大為原來的幾倍。

• 如何使DistributedDataParallel支持Sparse Embedding? 可以參考我在 PyTorch 論壇上的回答 DistributedDataParallel Sparse Embeddings [13]，設置torch.distributed.init_process_group中的backend=gloo即可，現在版本（1.6 以及 nightly）的 PyTorch 在nccl中仍然不支持Sparse Embedding。關于這個問題的最新進展可以看這個 PR：Sparse allreduce for ProcessGroupNCCL [14]

最后，PyTorch 1.6 中提供了混合精度訓練 amp [15] 的 API，可以方便的調用，通過將一部分操作從torch.float32 (float)變為torch.float16 (half)來加快訓練速度。但是它并不支持Sparse Embedding，如果你的模型中包含Sparse向量的話，最好不要使用。

參考文獻

[1] https://www.dgl.ai/

[2] https://pytorch.org/

[3]?https://docs.dgl.ai/guide/minibatch.html#chapter-6-stochastic-training-on-large-graphs

[4]?https://github.com/dmlc/dgl/tree/master/examples/pytorch/GATNE-T

[5]?https://networkx.github.io/documentation/latest/reference/algorithms/generated/networkx.algorithms.components.connected_components.html#networkx.algorithms.components.connected_components

[6] https://discuss.pytorch.org/t/num-workers-in-dataloader-will-increase-memory-usage/28522

[7] https://github.com/pytorch/pytorch/issues/13246

[8]?https://github.com/dmlc/dgl/blob/master/examples/pytorch/pinsage/model_sparse.py

[9] https://pytorch.org/docs/stable/optim.html?highlight=sparsea#torch.optim.SparseAdam

[10]?https://docs.python.org/3/library/multiprocessing.html

[11]?https://github.com/dmlc/dgl/blob/master/examples/pytorch/graphsage/train_sampling_multi_gpu.py

[12] https://discuss.pytorch.org/t/should-we-split-batch-size-according-to-ngpu-per-node-when-distributeddataparallel/72769

[13]?https://discuss.pytorch.org/t/distributeddataparallel-sparse-embeddings/60410/2

[14] https://github.com/pytorch/pytorch/issues/22400

[15] https://pytorch.org/docs/stable/amp.html

更多閱讀

#投 稿?通 道#

?讓你的論文被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。

?????來稿標準：

? 稿件確系個人原創作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?

? 如果文章并非首發，請在投稿時提醒并附上所有已發布鏈接?

? PaperWeekly 默認每篇文章都是首發，均會添加“原創”標志

?????投稿郵箱：

? 投稿郵箱：hr@paperweekly.site?

? 所有文章配圖，請單獨在附件中發送?

? 請留下即時聯系方式（微信或手機），以便我們在編輯發布時和作者溝通

????

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

關于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結

以上是生活随笔為你收集整理的大规模图训练调优指南的全部內容，希望文章能夠幫你解決所遇到的問題。

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