一、訓練大規模圖
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% 左右就可以了,不然反而可能會因為開執行緒的消耗拖慢訓練速度)。
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
向量的話,最好不要使用。
參考文獻