02.26 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀




GNN在社交網絡、知識圖、推薦系統各個領域得到了越來越廣泛的應用。19年的時候,聽facebook那邊的童鞋分享,第一次瞭解到BigGraph在工業界的應用。BigGraph是GNN在工業應用中一個比較優秀的實踐,從而萌生了我在實際工作場景中實際應用的想法。在實際的工作場景中和對BigGraph的二次開發過程中,對BigGraph調度分佈式模式的理解成為應用場景中比較關鍵的一環。本文是基於BigGraph的官方文檔翻譯來的,加上了一些我們實踐過程中對系統的理解。能力有限,如有問題,歡迎指正。


圖嵌入是一種從圖中生成無監督節點特徵(node features)的方法,生成的特徵可以應用在各類機器學習任務上。現代的圖網絡,尤其是在工業應用中,通常會包含數十億的節點(node)和數萬億的邊(edge)。這已經超出了已知嵌入系統的處理能力。我們介紹了一種嵌入系統,PyTorch-BigGraph(PBG),系統對傳統的多關係嵌入系統做了幾處修改讓系統能擴展到能處理數十億節點和數萬億條邊的圖形。

PBG使用了圖形分區來支持任意大小的嵌入在單機或者分佈式環境中訓練。我們展示了在功能的基準上與現有嵌入系統的性能比較,同時PBG允許在多臺機器上允許縮放到任意大小並且支持並行。我們在使用幾個大型社交圖網絡作為完整Freebase數據集來訓練和評估嵌入系統,數據集包含超過1億個節點和20億條邊。


簡介

圖結構數據是各種機器學習任務的一種常見輸入,直接處理圖結構數據是比較困難的,常用的技術是通過圖嵌入方法為圖中的每個節點創建向量化表示使得這些向量間的距離能預測圖形中是否存在邊。圖嵌入已經被證明是下游任務中有意義的特徵,如:電子商務中的推薦系統,社交媒體中的鏈接預測,藥物相互作用、表徵蛋白網絡。


圖結構數據在現代網絡公司非常常見,這給標準的嵌入方法提出了額外的調整:規模。例如:Facebook圖中包含20億個用戶節點和超過1萬億條邊,這些邊代表朋友關係、喜好、帖子和其他鏈接。阿里巴巴的用戶和產品圖也包含10億以上的用戶和20億以上的商品。在Pinterest,用戶到項目的圖包含20億的實體和超過170億邊。對這樣大小的圖做嵌入主要有兩個挑戰,一是這個嵌入系統必須足夠快,需要在一個合理的時間內完成10^11 - 10^12條邊的嵌入;二是擁有20億個節點,每個節點32個嵌入參數(浮點表示),大概需要800GB大小的內存來存儲這些參數,因此許多標準方法超出了典型商用服務器的內存容量。


我們介紹了Pytorch(譯註:facebook開源的深度學習框架)Biggraph,一種基於標準模型做了若干改進的嵌入系統。PBG帶來的改進是講標準方案擴展到具有數十億個節點和數萬億個邊圖。


PBG的重要組成部分:

1、將鄰接矩陣分塊,分解為n個桶,每次從一個桶的邊緣開始進行訓練,然後PBG要麼交換每個分區到磁盤的嵌入來減少內存的使用,要麼跨機器進行分佈式訓練。


2、分佈式執行模型:對大參數矩陣進行快分解,以及用於特徵化節點的全局參數和特徵嵌入的參數服務架構


3、高效的節點負採樣:在數據中對負向節點進行均勻採樣 並在批處理中重用負節點以減少內存用量


4、支持多實體、多關係圖,支持邊緣權重、關係運算符選擇等關係配置選項


我們在Freebase、LiveJournal和YouTube圖數據集上評估了PBG,表明了PBG能和當前已有的陷入系統性能能匹配。


我們也給出了更大圖數據上的結果,我們構造了一個完整的Freebase圖(1.21億個實體,24億條邊)的嵌入,並將其公開發布。Freebase圖的劃分減少了88%的內存消耗,嵌入向量的質量只有略微的下降,在8臺機器上分佈式執行,時間減少了4倍。另外我們在一個Twitter的大圖數據上進行了試驗,通過近似線性縮放得到了相似的結果。


PBG 作為一個開源項目發佈在:

https://github.com/facebookresearch/PyTorch-BigGraph.PBG是通過Pytorch實現的,沒有其他外部依賴。


BigGraph 中文文檔

本系列為翻譯的pytouch的官方手冊,希望能幫助大家快速入門GNN及其使用,全文十五篇,文中如果有勘誤請隨時聯繫。


1.數據模型

源鏈接:

https://torchbiggraph.readthedocs.io/en/latest/data_model.html

PBG operates 是有向多關係多圖,圖上的頂點稱為實體,每個邊將源連接到目標實體。源和目標分別稱為左側和右側(簡稱lhs和rhs)。同一對實體之間允許有多個邊,也允許使用環,即左側和右側相同的邊。


每個實體都屬於特定的實體類型(即每個實體只有一種類型)。所以類型將所有實體劃分為不相交的組。同樣,每個邊也只屬於一個關係類型,並且關係類型的邊的左側實體應該有相同實體類型,以及右側應該也是相同實體類型(左右側實體類型可以具有不同的實體類型)。此屬性表示每個關係類型都有左側實體類型和右側實體類型。

AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

在這個圖中,有14個實體:5個紅色實體、6個黃色實體、3個藍色實體;同時,圖中有12條邊,6條橘色的邊,3條紫色的邊,3條綠色的邊


為了讓PBG能夠讓大規模的圖形能運行,我們將圖形分解成小塊,在這些小塊上可以進行分佈式的訓練。首先是通過將每種類型的實體進一步拆分為若干個子集(稱為分區)。然後,對於每種關係類型,它的邊被劃分為桶:對於每對分區(一個來自左側,另一個來自右側實體類型)將創建一個桶,其中包含特定類型的邊及對應左、右側實體。

AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

本圖展示了一種實體劃分的方式,將紅色實體分為3類、黃色分為3類,藍色僅為1類。邊顯示的是黃色實體分桶1和紅色實體分桶2之間的實體邊。


注:出於技術原因,當前出現在關係類型左側的所有實體類型必須劃分為相同數量的分區(未分區的實體除外)。對於出現在右側的所有實體類型,必須保持相同的狀態。在numpy中,這意味著所有實體的分區數必須可以擴展到相同的值。


一個實體通過自身的類型、分桶、分組中的索引來定義(索引必須是連續的,這意味著如果一個類型的分區中有n個實體,則其索引位於半開放區間[0,n])。一條邊銅鼓變得類型,分桶(即其左側和右側實體類型的分區)以及其左側和右側實體在各自分區中的索引來標識。邊緣不必指定其左側和右側實體類型,因為它們隱式地存在於邊緣的關係類型中。


形式上,每個桶都可以通過一對整數(i,j)來標識,其中i和j分別是左側和右側分區。在該桶中,每個邊都可以由整數(x,r,y)的三聯體標識,x和y分別表示左側和右側實體,r表示關係類型。通過首先在配置中查找關係類型R來“解釋”此邊緣,並發現其左側只能有類型e1的實體,右側只能有類型e2的實體。然後可以確定左側實體,由(e1,i,x)(其類型、分區及其在分區內的索引)給出,同樣,也可以確定右側實體(e2,j,y)。


2.從實體嵌入到邊分值

源鏈接:

https://torchbiggraph.readthedocs.io/en/latest/scoring.html


Embeding 訓練的目標是將每個實體嵌入到R空間中,並且讓兩個實體的嵌入能很好地預測他們之間是否存在某種特定的關係。


更準確地說,訓練的目標是為每個實體學習嵌入併為每個關係類型學習一個函數。函數接受兩個實體嵌入,為它們分配一個分數。函數主要是為了讓相關實體比不相關實體能獲得更高的分數。


訓練集中提供的所有邊都被視為正實例。為了進行訓練,還需要一組負邊。這些負邊不是由用戶提供的,而是由系統在培訓期間生成的(參見見負抽樣),通常通過固定左側實體和關係類型並抽樣一個新的右側實體,反之亦然。這種抽樣方案在大型稀疏圖是合理的,因為這樣生成的邊在圖中是真正邊的概率很低。


通常來說,實體嵌入可以在R空間中取任何值。但是,在某些情況下(例如,當限制它們在某個球內時,或者當使用餘弦距離比較它們時),它們的“角度”將比它們的範數更重要。


每個關係評分函數必須以特定的形式表示(文獻中最常見的功能可以轉換為這種表示)。在當前的實現中,只允許對其中一個邊的嵌入進行轉換,然後使用通用的對稱比較器函數將其與另一個邊的未轉換嵌入進行比較,這對於所有關係都是相同的。形式上,對於左側實體x和右側實體y,以及對於關係類型r,得分為:fr(θx,θy)=c(θx,gr(θy)).其中θx和θy分別是x和y的嵌入,fr是r的得分函數,gr是r的操作器,c是比較器。


在“正常”情況下(所謂的“標準”關係模式),運算符僅應用於右側實體。使用動態關係時不是這樣。將運算符應用於兩側通常是多餘的。另外,與另一側相比,首選一側可以打破對稱性並捕捉邊緣的方向。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Embeddings 嵌入

嵌入表示存在在D維真實空間中,其中D由維度配置參數確定。通常,每個實體都有自己的嵌入,這完全獨立於任何其他實體的嵌入。當使用特徵化實體的方法不同,是實體的嵌入將是其特徵嵌入的平均值。


如果設置了max_norm配置參數,則每次參數更新後,嵌入將以半徑max_norm投影到單位球上。如果要添加新的嵌入類型,需要實現torchbigraph.model.abstractEmbedding類的子類。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Global embeddings 全局嵌入

當global_emb配置參數被激活,每個實體嵌入通過一個特定的實體類型向量(與嵌入同時學習)來轉化。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Operators 操作器

none,無操作,使嵌入保持不變;

translation, 平移,增加了嵌入相同維的向量;

diagonal, 對角線,將每個尺寸乘以不同的係數(相當於乘以對角線矩陣);

linear, 線性,應用線性映射,即乘以一個全平方矩陣。

affine, 仿射,應用仿射變換,即線性後接translation。

complex_diagonal, 複數_對角線,它將D-維實向量解釋為D/2維復向量(D必須是偶數;向量的前半部分是實向量,後半部分是虛向量),然後將每個條目乘以不同的複雜參數,就像diagonal一樣。

所有操作器的參數是在訓練中學得的。如果要自定義操作器,我們需要實現torchbiggraph.model.AbstractOperator的子類(在動態關係情況下實現torchbiggraph.model.AbstractDynamicOperator子類,docstrings解釋了必須實現什麼)並且在torchbiggraph.model.register_operator_as()裝飾器中註冊(或者torchbiggraph.model.register_dynamic_operator_as())指定一個新名稱,然後在配置中使用該名稱來選擇比較器。上述所有操作都可以在配置文件內部完成。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Comparators 比較器

可用的比較器有:

dot,點積,計算兩個嵌入向量的標量或內積;

cos,cos距離,是兩個向量之間夾角的餘弦,或等於點積除以向量範數的乘積。

l2,負的l2距離,也就是歐幾里得距離(負是因為較小的距離會得到較高的分數)。

squared_l2,負平方L2距離。

自定義比較器需要實現torchbiggraph.model.AbstractComparator 子類並且在torchbiggraph.model.register_comparator_as()裝飾器中註冊,指定一個新名稱,然後在配置中使用該名稱來選擇比較器。上述所有操作都可以在配置文件內部完成。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Bias 偏置

如果使用了bias 配置,那麼嵌入的第一個座標將作為比較器計算中的偏差。這意味著比較器將僅根據向量的最後一個D-1條目進行計算,然後將兩個向量的第一個條目都添加到結果中。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Coherent sets of configuration parameters 相關配置集

本章中描述的參數在配置文件中集中顯示(為了更接近實現,並允許更靈活的調優),但它們的某些組合比其他組合更優。

除了默認配置外,還發現以下的配置效果較優:init_scale = 0.1, comparator = dot, bias = true, loss_fn = logistic, lr = 0.1.


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Interpreting the scores 分值解釋

在訓練過程中,不同的應用會基於不同的損失函數,這讓得到的分值有不同的解釋。常見的包括對和實體有關聯關係的實體排序,確定兩個給定實體之間存在某種關係的概率等。


3.I/O格式化

原鏈接:

https://torchbiggraph.readthedocs.io/en/latest/input_output.html


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Entity and relation types 實體和關係類型

配置文件中需要包含實體字典,文件內容是實體類型(每個類型都由字符串標識)以及別的有關的信息。配置文件的relations key是關係類型的列表(每個類型都由其在該列表中的索引標識)以及別的數據如左右兩側的實體類型。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Entities 實體

需要提供的關於實體的唯一信息是每個實體類型的分區中有多少個實體。這是通過將名為entity_count_type_part.txt的文件放在entity_path config參數指定的目錄中,用於由類型標識的每個實體類型和每個分區。這些文件必須包含一個整數(文本),即該分區中的實體數。必須將所有這些文件所在的目錄指定為配置文件的entity_path路徑。


我們可以通過在init_path配置項中設置特定的值來為embeddings初始化,配置項的值為包含文件內容格式類似於Checkpoint文件的路徑,輸出格式在Checkpoint中詳細介紹。


如果沒有提供初始值,每個維度都從中心正態分佈中取樣來自動生成,標準偏差可以使用init_scale配項鍵進行配置。出於性能原因,特定類型的所有實體的樣本將不獨立。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Edges邊

對於每個bucket必須有一個文件來存儲屬於該bucket的所有邊和所有關係類型。這代表文件只能由兩個整數來標識,即其左側和右側實體的分區。它必須命名為edges_lhs_rhs.h5(其中lhs和rhs是上面的整數),文件必須是hdf5格式,包含三個長度相同的一維數據集,稱為rel、lhs和rhs。每個元素中第i個位置的元素定義了第i個邊緣:rel標識關係類型(從而標識左側和右側實體類型),lhs和rhs給出了各自分區中左側和右側實體的索引。


為了方便將來格式更新,每個文件需要在頂級組的format_version屬性中包含格式版本,當前版本為1。


如果一個實體類型是未分區的(即所有實體都屬於同一分區),那麼這些實體的邊必須均勻分佈在所有存儲桶中。


這些文件,對於所有的桶,必須存儲在同一目錄中,該目錄通過edge_paths來配置。這個配置實際上可以包含一個路徑列表,每個路徑指向上面描述的格式的目錄:如果這樣配置圖形將包含所有邊的並集。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Checkpoint 檢查點

訓練的數據,也就是檢查點被寫入在checkpoint_path中配置的目錄中。checkpoint 由連續的正整數(從1開始)標識,並且屬於某個checkpoint 的所有文件在其名稱和擴展名之間都有一個額外的組件.vversion(例如,對於版本42為.v42.h5)。


最新的完整checkpoint版本存儲在同目錄下名為checkpoint_version.txt文件中,其中包含一個整數,即當前版本。


每個checkpoint都包含一個配置的JSON轉儲文件,用於存儲在config.json文件中的配置


保存新版本checkpoint後,將自動刪除以前的版本。為了定期保留其中一些版本,請將checkpoint_preservation_interval 配置設置為所需的週期(以epoch數表示)。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Model parameters 模型參數

模型參數存儲在名為model.h5的文件中,文件格式是一個hdf5文件,其中包含每個參數的一個數據集,所有數據集都位於model 組中。目前提供的參數有:

model/relations/idx/operator/side/param 存儲了每個操作器的關係參數

model/entities/type/global_embedding 存儲了每個預訓實體類型的全局嵌入


這些數據集中的每一個還包含存儲在模型狀態字典,用state-dict-key屬性配置。另外還可能有另一個數據集optimizer/state_dict,包含模型優化器狀態字典的二進制文件(可以通過torch.save()獲取)。


最後,文件的頂級組包含包含了一些附加元數據的屬性。主要包括格式版本、配置的JSON轉儲和一些有關生checkpoint的迭代的信息。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Embeddings 嵌入

對於每個實體類型及其分區,都有一個文件embeddings_type_part.h5(其中type是類型的名稱,part是分區的基於0的索引)。該文件是包含兩個數據集的HDF5文件——嵌入,包含實體的嵌入:一個二維數據集,第一個維度是實體數,第二個維度是嵌入的維度。


與模型參數文件一樣,優化器狀態dict和其他元數據也包含在裡面。


圖嵌入是一種從圖中生成無監督節點特徵(node features)的方法,生成的特徵可以應用在各類機器學習任務上。現代的圖網絡,尤其是在工業應用中,通常會包含數十億的節點(node)和數萬億的邊(edge)。這已經超出了已知嵌入系統的處理能力。我們介紹了一種嵌入系統,PyTorch-BigGraph(PBG),系統對傳統的多關係嵌入系統做了幾處修改讓系統能擴展到能處理數十億節點和數萬億條邊的圖形。


4.批預處理

原鏈接:

https://torchbiggraph.readthedocs.io/en/latest/batch_preparation.html


本節介紹了每一批次數據在損失的計算和優化前是如何準備和組織的。


訓練通過循環嵌套來迭代處理邊。掃描從外層到內層叫代(時期),代際間是相互獨立並且過程基本一致,這樣的目的是重複內部循環直到收斂。每代迭代會訪問到所有的邊,epochs的數目在num_epochs配置參數中指定。


譯者注:

(1)batch:1個batch代表深度學習算法一次參數的更新,所需要損失函數並不是由一個數據獲得的,而是由一組數據加權得到的,這一組數據的數量叫batchsize;

(2)iteration:1個iteration等於使用batchsize個樣本訓練一次;

(3)epoch:1個epoch等於使用訓練集中的全部樣本訓練一次,通俗的講epoch的值num_epoch就是整個數據集被輪幾次。


當迭代一個邊集合時,每個分桶首先會被分割為大小相等的塊:每個塊由連續間隔的邊構成(和存儲在文件中的順序一致),塊的數量可以通過設置num_edge_chunks來調整。訓練從所有桶的第一個塊開始,然後第二個,以此類推。


然後算法開始迭代桶,訓練桶的順序依賴於配置文件中bucket_order這個配置項。除了隨機排列之外,有一些方法用於嘗試將訓連續的桶之間共享一個分區:這讓分區能夠被重用,從而允許可以允許參數存儲在內存中而不是先註銷然後被另一個桶加載到自己的空間中(在分佈式模式下,不同的訓練器進程同時在桶上運行,因此迭代的管理方式不同)。


當訓練器被固定了一個確定的桶和一個指定的快,邊會最終磁盤上加載起來。當在訓練中評估的時候,這些邊的一個子集被保留(這個子集對所有的代都是一樣的)。被保留的這些邊會被均勻打散然後分為相等的部分。為了讓訓練過程中可以訓練過程可以在同一時間並行進行,打散後的分片會放到一個分佈式的處理池中處理Processes。這些子流程相互間不同步他們的計算或內存,我們叫 “Hogwild” worker,這些節點worker的數據量通過workers參數來控制。


每個worker上訓練邊集合的方式取決於動態關係是否被用了。最簡單的情況是邊被分割到連續的batches(每個batch都和batch_size中配置指定的大小一樣,除了最後一個可以略小一點),這樣便可以對樣本一批接一批順序訓練。


當動態關係沒有被使用,也就是說損失的計算只能對一個邊的集合中所有具有相同類型的邊計算。因此,worker首先隨機一個關係類型,選取的數量應該和樣本池中同類型邊的概率佔比一致。然後採用手個batch_size大小的同類關係(如果剩下的樣本不夠),在訓練池中刪除掉這些樣本並開始訓練。


5.分佈式模式

源鏈接:

https://torchbiggraph.readthedocs.io/en/latest/distributed_training.html


PBG可以跨多臺機器訓練,通過網絡進行機器間通信,以減少大型圖形上的訓練時長。分佈式訓練可以同時利用更大的計算資源,同時將整個模型跨機器存儲在內存中,避免了內存和硬盤間的數據交換。訓練在每個機器上多個子流程間進一步的並行化。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Setup 啟動

若要執行分佈式的訓練,我們首先得更新包含在分佈式配置中的必須的一些配置項。如果訓練要在N個機器上執行,那在配置中的num_machines配置項應該要正確配置。另外,訓練機器相互之間發現和通信模式的設置通過distributed_init_method來描述。所有torch.distributed.init_process_group()可接收的參數都是init_methed有效參數。通常這是一個共享網絡文件系統的路徑或其中一臺計算機的網絡地址。相關的信息和完整參考,請參閱pytorch文檔。


啟動分佈式訓練,在每臺機器上啟動torchbiggraph_train --rank rank config.py ,其中對每臺機器,rank需要替換成從0到N-1的不同的整數。每個機器上必須已經安裝了PBG並且都有一份相同的配置文件。


在某些非常規的情況下,需求是希望將嵌入存放到不同的機器上而不是訓練嵌入。在這種情況下,可以配置num_partition_servers及存放的機器數,並且在一些實例上啟動torchbiggraph_partitionserver。請參閱下面的詳細信息。


Tip:

推薦的默認設置是將num_machines設置為分區數量的一半(參見下面的原因),並讓num_partition_servers保持未設置狀態。


一旦啟動所有這些命令,就不再需要手動干預。


Warning:

未分區的實體不應該用於做分佈式訓練。原因是儘管分區實體類型的嵌入一次只能在一臺計算機上使用,並可以根據需要在計算機之間進行交換,但未分區實體類型的嵌入是通過一個優化不良的參數服務器異步通信的,該服務器是為共享關係而設計的,空間很小,所以不支持同步大量參數。例如,具有1000多個實體的未分區實體類型,在這種情況下,未劃分的嵌入的質量可能非常差。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Communication protocols 通信協議

分佈式訓練包含機器間協同和通信的在不同目標下的不同方式,包含以下幾種:


1)廣播那個訓練器正在操作哪個分桶,通過分配來避免衝突;

2)在需要時將實體分區的嵌入從一個訓練器傳遞到下一個訓練器(針對一次只能由一個訓練器訪問的數據類型)

3)通過收集和重新分發更新參數,共享所有訓練器需要同時訪問的參數。


上述的方案都被獨立實現為“協議”,每個訓練器通過啟動子流程作為協議的一部分或者全部,在不同的協議中扮演客戶端或者服務端的角色。下面將對這些協議進行說明,以深入瞭解系統。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Synchronizing bucket access 同步分桶訪問

PBG通過在多機上訓練不交叉的分桶來實現分佈式訓練(如:分桶間不會包含有相同的partition)。因此,每個分區一次最多由一臺機器使用,每臺機器最多使用兩個分區(唯一的例外是“對角線上”的存儲桶,它們具有相同的左側和右側分區)。這意味著可以同時訓練的bucket數量大約是分區總數的一半。


這種方式讓集群通過“lock server”商定哪個機器可以對哪個桶進行操作。服務器隱式的從排序為0訓練器開始,其他所有機器作為他的客戶端,請求一個新的分桶來開始訓練操作(當他們需要新的分桶時),獲得一個服務器指定的分桶(或者未獲得,如果所有的分桶都已經被訓練或者在“locked”狀態,即該partitions被其他訓練器使用中),開始訓練,訓練完成後返回並重覆上述動作。lock server 優化I/O的策略:當訓練器請求bucket時,儘可能的分配和上一個bucket相鄰的bucket到的相同的訓練器上,這樣這些分區就可以保存在內存中,而不必卸載和重新加載。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

Exchanging partition embeddings 分區交換嵌入

當一個訓練器開始對一個分桶的數據操作的時候,它需要訪問所有實體的嵌入向量embeddings,這些向量不然是分桶的左鄰接節點,不然是右鄰接節點。locking 機制通過lock server確保在給定時刻內只有一個訓練器對partition進行操作。但是這無法處理未分區的實體類型,這需要再所有訓練器中共享,參照下文。因此保證了每個訓練器對自己持有的partition是獨佔的。


一旦一個訓練器開始在新的bucket上工作,就需要獲取其分區的嵌入,在完成後釋放,並在其更新版本中提供給下一個需要它們的訓練器。為了支持上述過程,有一個顧名思義叫“partition servers”的系統來存儲嵌入向量,提供給需要這些嵌入的訓練器並接收對這些向量的更新並存儲。


這個服務是可選的,當吧num_partition_servers設置為0時被禁用。在這種情況下,訓練器只需將嵌入內容寫入檢查點目錄(該目錄應位於共享磁盤上),然後其他訓練器從那裡取回,就可以相互“傳送”。


啟用此係統後,它可以在兩種模式下工作。最簡單的模式是我們把num_partition_servers設置為-1(默認):在這種情況下,所有訓練器會生成一個本地進程作為partition服務。另外一種模式則num_partition_servers是一個正值,那麼訓練器將不會生成任何進程,但是我們需要通過在適當數量的機器上手動通過torchbiggraph_partitionserver命令來啟動該系統。


AI | 圖神經網絡-Pytorch Biggraph簡介及官方文檔解讀

共享參數更新

有些模型參數需要一直被全局的訓練器訪問(包含操作器operator的權重,每個實體類型的全局嵌入,未分區實體的嵌入向量)。有些參數不依賴於訓練器處理的分桶,但所有訓練器都需要常駐以使用到(和跟進需要反覆加載和卸載的實體嵌入相反)。這些參數通過一組“parameter servers”來同步,每個訓練器啟動一個本地的參數服務器(一個獨立的子進程)並且和其他參數服務器連接。訓練器間共享每個參數並且存儲在parameter server(如果太大的話可能是分片)。每個訓練器也有一個循環(也在一個單獨的子進程中)定期遍歷每個共享參數,計算其當前本地值和上次與參數的服務器同步時的值之間的差值,並將該差值發送給服務器。服務器依次累加從所有訓練器接收到的所有增量,更新參數值並將此新值發送回訓練器。同時,為了防止參數服務器使其他通信不足,參數服務器限制100個更新/s或1Gb/s。


參考文獻


[1]https://github.com/facebookresearch/PyTorch-BigGraph.PBG

[2] https://www.jianshu.com/u/a0aed87ff1be


分享到:


相關文章: