谷歌正式發(fā)布TensorFlow 圖神經(jīng)網(wǎng)絡(luò)

發(fā)布人：Sibon Li、Jan Pfeifer、Bryan Perozzi 和 Douglas Yarrington

日前，我們很高興發(fā)布了 TensorFlow 圖神經(jīng)網(wǎng)絡(luò) (Graph Neural Networks, GNNs)，此庫(kù)可以幫助開(kāi)發(fā)者利用 TensorFlow 輕松處理圖結(jié)構(gòu)化數(shù)據(jù)。我們已在 Google 的多個(gè)生產(chǎn)環(huán)境中使用了該庫(kù)的早期版本（例如，垃圾郵件和異常檢測(cè)、數(shù)據(jù)流量估計(jì)、YouTube 內(nèi)容標(biāo)簽），并作為可擴(kuò)容的圖挖掘管道的一個(gè)組成部分。特別是，鑒于 Google 的數(shù)據(jù)類(lèi)型繁多，我們的庫(kù)在設(shè)計(jì)時(shí)就考慮到了異構(gòu)圖。發(fā)布此庫(kù)的初衷是為了鼓勵(lì)與業(yè)界研究人員的合作。

TensorFlow 圖神經(jīng)網(wǎng)絡(luò)

https://github.com/tensorflow/gnn

為何使用 GNNs？

在現(xiàn)實(shí)世界和我們的工程系統(tǒng)中，“圖”無(wú)處不在。一組物體、地點(diǎn)或人以及它們之間的聯(lián)系通常都可以用圖來(lái)表述。通常情況下，我們?cè)?a target="_blank">機(jī)器學(xué)習(xí)問(wèn)題中看到的數(shù)據(jù)是結(jié)構(gòu)化或關(guān)系化的，因此也可以用圖來(lái)表述。雖然關(guān)于 GNNs 的基礎(chǔ)研究可能只有幾十年的歷史，但當(dāng)代 GNNs 最近取得的功能進(jìn)展已經(jīng)幫助推動(dòng)了多個(gè)領(lǐng)域的進(jìn)步，其中包括數(shù)據(jù)流量預(yù)測(cè)、謠言和假新聞檢測(cè)、疾病傳播建模、物理學(xué)模擬和理解分子有氣味的原因等。

數(shù)據(jù)流量預(yù)測(cè)

https://deepmind.com/blog/article/traffic-prediction-with-advanced-graph-neural-networks

謠言和假新聞檢測(cè)

https://arxiv.org/abs/2108.03548

疾病傳播建模

https://arxiv.org/abs/2007.03113

物理學(xué)模擬

http://proceedings.mlr.press/v80/sanchez-gonzalez18a/sanchez-gonzalez18a.pdf

理解分子有氣味的原因

https://arxiv.org/abs/1910.10685

圖可以對(duì)多種不同類(lèi)型數(shù)據(jù)之間的關(guān)系進(jìn)行建模，

包括網(wǎng)頁(yè)（左）、社交關(guān)系（中）或分子（右）等

圖代表了一組實(shí)體（節(jié)點(diǎn)或頂點(diǎn)）之間的關(guān)系（邊）。我們可以通過(guò)描述每個(gè)節(jié)點(diǎn)、邊或整個(gè)圖的特征，從而將信息存儲(chǔ)在圖的每一塊中。此外，我們還可以賦予邊方向性，來(lái)描述信息或數(shù)據(jù)流，等等。

若這些圖有多個(gè)特征，則可以用 GNNs 來(lái)解決這類(lèi)問(wèn)題。通過(guò)圖層級(jí)的研究，我們嘗試預(yù)測(cè)整個(gè)圖的特征。我們可以識(shí)別某些“形狀”的存在，如圖中的圓圈，可能代表亞分子，也可能代表密切的社會(huì)關(guān)系。GNNs 可以用于節(jié)點(diǎn)級(jí)的任務(wù)，對(duì)圖的節(jié)點(diǎn)進(jìn)行分類(lèi)，并預(yù)測(cè)圖中的分區(qū)和相似性，類(lèi)似于圖像分類(lèi)或分割。最后，我們可以在邊層級(jí)上使用 GNNs 來(lái)發(fā)現(xiàn)實(shí)體之間的聯(lián)系，或許可以使用 GNNs 來(lái)“修剪”邊，以確定場(chǎng)景中對(duì)象的狀態(tài)。

結(jié)構(gòu)

TF-GNN 提供了基本模塊，以便在 TensorFlow 中實(shí)現(xiàn) GNN 模型。除了建模 API，我們的庫(kù)還針對(duì)處理圖數(shù)據(jù)的困難任務(wù)提供了大量的工具：基于 Tensor 的圖數(shù)據(jù)結(jié)構(gòu)、數(shù)據(jù)處理管道，以及一些供用戶(hù)快速上手的示例模型。

組成工作流的各種 TF-GNN 組件

TF-GNN庫(kù)的初始版本包含一些實(shí)用程序和功能，初學(xué)者和有經(jīng)驗(yàn)的用戶(hù)都能使用，其中包括：

TF-GNN 庫(kù)

https://github.com/tensorflow/gnn

高階 Keras 式 API，用于創(chuàng)建 GNN 模型，可以輕松地與其他類(lèi)型的模型組合。GNNs 經(jīng)常與排名、深度檢索（雙編碼器）結(jié)合使用或與其他類(lèi)型的模型（圖像、文本等）混合使用。

用于異構(gòu)圖的 GNN API。我們?cè)?Google 和現(xiàn)實(shí)世界中處理的許多圖問(wèn)題都包含不同類(lèi)型的節(jié)點(diǎn)和邊。因此，我們選擇提供一種簡(jiǎn)單的方法來(lái)對(duì)此建模。

定義明確的架構(gòu)，用于聲明圖的拓?fù)?，以及?yàn)證架構(gòu)的工具。此架構(gòu)描述了其訓(xùn)練數(shù)據(jù)的形狀，并用于指導(dǎo)其他工具。

GraphTensor 復(fù)合張量類(lèi)型，它持有圖數(shù)據(jù)，可以分批處理，還有圖操作例程可用。

關(guān)于 GraphTensor 結(jié)構(gòu)的操作庫(kù)：

對(duì)節(jié)點(diǎn)和邊的各種有效的廣播和池化運(yùn)算，以及相關(guān)工具。

標(biāo)準(zhǔn)并入的卷積庫(kù)，ML 工程師/研究人員可以輕松地對(duì)其進(jìn)行擴(kuò)展。

高階 API，幫助產(chǎn)品工程師快速構(gòu)建 GNN 模型，而不必?fù)?dān)心其細(xì)節(jié)問(wèn)題。

磁盤(pán)上的圖形訓(xùn)練數(shù)據(jù)的編碼，以及用來(lái)將該數(shù)據(jù)解析為數(shù)據(jù)結(jié)構(gòu)的庫(kù)，您的模型可以從該庫(kù)中提取各種特征。

示例用法

在下面的例子中，我們使用 TF-GNN Keras API 構(gòu)建一個(gè)模型，根據(jù)用戶(hù)觀看的內(nèi)容和喜歡的類(lèi)型向其推薦電影。

我們使用 ConvGNNBuilder 方法來(lái)指定邊類(lèi)型和節(jié)點(diǎn)配置，即對(duì)邊使用 WeightedSumConvolution（定義如下）。每次通過(guò) GNN 時(shí)，我們將通過(guò) Dense 互連層來(lái)更新節(jié)點(diǎn)值：

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)
    ])

以上代碼很好用，但有時(shí)我們可能想要為 GNNs 使用更強(qiáng)大的自定義模型架構(gòu)。例如，在之前的用例中，我們可能想指定某些電影或類(lèi)型，讓它們?cè)谖覀冞M(jìn)行推薦時(shí)擁有更多權(quán)重。在下列片段中，我們用自定義圖卷積定義了一個(gè)更高級(jí)的 GNN，例子中使用的是加權(quán)邊。我們定義了 WeightedSumConvolution 類(lèi)來(lái)匯集邊值，讓其作為所有邊的權(quán)重之和：

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

請(qǐng)注意，盡管卷積的編寫(xiě)只考慮了源節(jié)點(diǎn)和目標(biāo)節(jié)點(diǎn)，但 TF-GNN 保障了其適用性，讓其能夠在異構(gòu)圖（有各種類(lèi)型的節(jié)點(diǎn)和邊）上無(wú)縫工作。

更多信息

您可以查看 TF-GNNGitHub repo，以獲得更多信息。

GitHub repo

https://github.com/tensorflow/gnn

歡迎加入 TensorFlow 論壇，隨時(shí)關(guān)注 TensorFlow 博客、Twitter 或 YouTube，獲悉最新消息。

TensorFlow 論壇

https://discuss.tensorflow.google.cn/

博客

https://blog.tensorflow.google.cn/

Twitter

http://twitter.com/tensorflow

YouTube

http://youtube.com/tensorflow

您還可以通過(guò) Community Spotlight 計(jì)劃向我們提交作品，分享您的構(gòu)建成果。如需提供反饋，您可以在 GitHub 上提交問(wèn)題。我們歡迎您的貢獻(xiàn)和參與，謝謝！

Community Spotlight 計(jì)劃

http://goo.gle/TFCS

GitHub

https://github.com/tensorflow/gnn/issues

致謝

本文所介紹的研究來(lái)自以下各位的合作成果：來(lái)自 Google 的 Oleksandr Ferludin?、Martin Blais、Jan Pfeifer?、Arno Eigenwillig、Dustin Zelle、Bryan Perozzi 和 Da-Cheng Juan，以及來(lái)自 DeepMind 的 Sibon Li、Alvaro Sanchez-Gonzalez、Peter Battaglia、Kevin Villela、Jennifer She 和 David Wong。

原文標(biāo)題：推出 TensorFlow 圖神經(jīng)網(wǎng)絡(luò) (GNNs)

文章出處：【微信公眾號(hào)：谷歌開(kāi)發(fā)者】歡迎添加關(guān)注！文章轉(zhuǎn)載請(qǐng)注明出處。