高級API、異構圖:谷歌發布TF-GNN,在TensorFlow創建圖神經網絡
今天,TensorFlow 官方博客發布了 TensorFlow Graph Neural Networks(TensorFlow GNN)庫 ,這個庫使得用戶在使用 TensorFlow 時能夠輕松處理圖結構數據。
此前,TensorFlow GNN 的早期版本已經在谷歌的各種應用中使用,包括垃圾郵件和異常檢測、流量估計、YouTube 內容標記等。特別是,考慮到谷歌數據種類繁多,該庫在設計時就考慮到了異構圖。

項目地址:https://github.com/tensorflow/gnn
為何使用 GNN?
無論是在現實世界中,還是在我們設計的系統中,圖無處不在。一組對象或是不同的人以及他們之間的聯系,通常可以用圖來描述。通常情況下,機器學習中的數據是結構化或關系型的,因此也可以用圖來描述。雖然 GNN 的基礎研究已經有幾十年的歷史,但近幾年才取得一些進展,包括在交通預測、假新聞檢測、疾病傳播建模、物理模擬,以及理解為什么分子會有氣味等。

圖可以為不同類型的數據進行關系建模,包括網頁(左)、社交關系(中)或分子(右)。
怎樣定義圖呢?簡單來講,圖表示一組實體(節點或頂點)之間的關系(邊)。我們可以描述每個節點、邊或整個圖,從而將信息存儲在圖的每一部分中。此外,我們可以賦予圖邊緣方向性來描述信息或信息流。
GNN 可以用來回答關于這些圖的多個特征問題。GNN 可用于節點級任務,對圖的節點進行分類,并預測圖中的分區和相關性,類似于圖像分類或分割。最后,我們可以在邊緣級別使用 GNN 來發現實體之間的連接。
TensorFlow GNN
TF-GNN(TensorFlow GNN) 提供了在 TensorFlow 中實現 GNN 模型的構建塊。除了建模 API 之外,該庫還為處理圖數據提供了可用工具,包括基于張量的圖數據結構、數據處理 pipeline 和一些供用戶快速入門的示例模型。

TF-GNN 工作流程組件
TF-GNN 庫的初始版本包含許多實用程序和功能,供初學者和有經驗的用戶使用,包括:
- 高級 keras 風格的 API 用于創建 GNN 模型,可以很容易地與其他類型的模型組合。GNN 通常與排序、深度檢索結合使用或與其他類型的模型(圖像、文本等)混合使用;
- 定義良好的模式用來聲明圖拓撲結構,以及驗證工具。該模式描述了其訓練數據的大小,并用于指導其他工具;
- GraphTensor 復合張量類型,可以用來保存圖數據,也可以進行批處理,并具有可用的圖操作例程;
- GraphTensor 結構操作庫:在節點和邊緣上進行各種有效的 broadcast 和 pooling 操作,以及提供相關操作的工具;標準 baked 卷積庫,機器學習工程師、研究人員可以對其輕松擴展;高級 API 可以幫助工程師快速構建 GNN 模型而不必擔心細節;
- 模型可以從圖訓練數據編碼,以及用于將此數據解析為數據結構的庫中提取各種特征。
示例
下面示例使用 TF-GNN Keras API 構建了一個模型,該模型可以根據觀看內容和喜歡的類型向用戶推薦電影。
完成這項任務使用 ConvGNNBuilder 方法來指定邊的類型和節點配置,即對邊使用 WeightedSumConvolution(定義如下):
- import tensorflow as tf
- import tensorflow_gnn as tfgnn
- # Model hyper-parameters:
- h_dims = {'user': 256, 'movie': 64, 'genre': 128}
- # Model builder initialization:
- gnn = tfgnn.keras.ConvGNNBuilder(
- lambda edge_set_name: WeightedSumConvolution(),
- lambda node_set_name: tfgnn.keras.layers.NextStateFromConcat(
- tf.keras.layers.Dense(h_dims[node_set_name]))
- )
- # Two rounds of message passing to target node sets:
- model = tf.keras.models.Sequential([
- gnn.Convolve({'genre'}), # sends messages from movie to genre
- gnn.Convolve({'user'}), # sends messages from movie and genre to users
- tfgnn.keras.layers.Readout(node_set_name="user"),
- tf.keras.layers.Dense(1)
- ])
有時我們希望 GNN 性能更強大,例如,在上個示例中,我們可能希望模型在給出推薦電影時可以同時給出權重。下面代碼片段中定義了一個更高級的 GNN,它帶有自定義圖卷積,以及帶有權重邊。下面代碼定義了 WeightedSumConvolution 類可以將邊值池化為所有邊的權重總和:
- class WeightedSumConvolution(tf.keras.layers.Layer):
- """Weighted sum of source nodes states."""
- def call(self, graph: tfgnn.GraphTensor,
- edge_set_name: tfgnn.EdgeSetName) -> tfgnn.Field:
- messages = tfgnn.broadcast_node_to_edges(
- graph,
- edge_set_name,
- tfgnn.SOURCE,
- feature_name=tfgnn.DEFAULT_STATE_NAME)
- weights = graph.edge_sets[edge_set_name]['weight']
- weighted_messages = tf.expand_dims(weights, -1) * messages
- pooled_messages = tfgnn.pool_edges_to_node(
- graph,
- edge_set_name,
- tfgnn.TARGET,
- reduce_type='sum',
- feature_value=weighted_messages)
- return pooled_messages
請注意,即使卷積是在只考慮源節點和目標節點的情況下編寫的,TF-GNN 仍可確保它適用并可以無縫處理異構圖(具有各種類型的節點和邊)。
安裝
這是目前安裝 tensorflow_gnn 的唯一方法。強烈建議使用虛擬環境。
Clone tensorflow_gnn:
- $> git clone https://github.com/tensorflow/gnn.git tensorflow_gnn
安裝 TensorFlow:
- $> pip install tensorflow
安裝 Bazel:Bazel 需要構建包的源代碼。安裝步驟請參考:https://docs.bazel.build/versions/main/install.html
安裝 GraphViz:這個包使用 GraphViz 作為可視化工具,安裝因操作系統而異,例如 Ubuntu:
- $> sudo apt-get install graphviz graphviz-dev
安裝 tensorflow_gnn:
- $> cd tensorflow_gnn && python3 -m pip install .