まえがき

勉強会グループ「機械学習名古屋」では、論文を読み発表し合うという「研究会」を今年（2018年）の4月から行っています。その中で"Relational inductive biases, deep learning, and graph networks"という論文を読みました。それには「近々ライブラリ化して公開するよ」と書かれており、実際に、GraphNetsという名前で10月にDeepMindから公開されました。
この記事では、ライブラリの内容やどのようなものを扱うのかについて勉強を兼ねて紹介しようと思います。

注）この記事でグラフと呼ぶものは、「辺」や「頂点」からなる方のグラフです。

インストール

GraphNetsはTensorFlowのライブラリSonnetで作られたライブラリです。Pythonの環境によって少々違いますが、次でインストールします。

$ pip install graph_nets tensorflow(もしくはtensorflow-gpu) tensorflow_probability

理論

ニューラルネットでグラフを扱う

ディープラーニングというと、画像・自然言語・音声を扱うことが多いように思います。これらは、格子もしくは系列の形をしたデータです。しかし、人間が認識・推論することや自然界の現象はより複雑に互いに関係しながら成り立っています。これはグラフのようなもので表現すると良さそうです。

グラフは、

隣接した頂点の数が一定ではない
形状が多様である

という部分が画像・自然言語・音声と異なります。とはいえ、グラフを扱うニューラルネットの研究は以前からあります。例えばGCN(graph-based convolutional network)があります。3Dメッシュによる物体の認識や、化学の分野で分子の研究であったりに利用されています。

GCNをざっくり説明すると、

グラフの各頂点に特徴ベクトルが配置される
（隣のマスの代わりに）その頂点とその隣接する頂点の特徴ベクトルにチャンネル毎に重みをかけて足し合わる
それを次の層のその頂点に対応した特徴ベクトルとする
訓練で上記の重みを更新する

というものです。

GraphNetwork

GraphNetsで扱うグラフは、単なるグラフではなく、

頂点（V）
辺（E）
グラフ全体（u）

に特徴ベクトルが付与されたものです。"Relational inductive biases, deep learning, and graph networks"は、このような値付きのグラフに対して、計算を行い新しい値付きのグラフを算出する GraphNetwork（GN）というものを提案しました。この計算は次のような図で表されます。

辺の計算。その辺とその辺の端点と全体の特徴ベクトルから計算。
頂点の計算。その頂点への辺の新しい特徴ベクトルを集計（平均を取るとか和を取るなど）。集計したベクトルと全体の特徴ベクトルから、その頂点の特徴ベクトルを計算。
全体の計算。全辺と全頂点の新しい特徴ベクトルを集計する。集計したベクトルともとの全体のベクトルから、新しい全体の特徴ベクトルを計算。

GCNと比較すると、

頂点以外にも特徴ベクトルがある
有向グラフが想定されている
GCNで隣接する頂点の特徴ベクトルに重みをつける部分は辺の計算に相当する
GCNで中心の頂点へ集約する方法が足し合わせと決め打ちされていたのに対して、GNはそれに制限せず何らかの関数を用いる
全体に対してもベクトルを割り当て、それぞれの計算で参照する

という違いがあります。GNの方がより広い概念になっています。

GNを何個も重ねたり、ときにリカレントな構造にするなどの自由度の高い利用方法が考えられます。GraphNetsはGNや関連するアルゴリズムを容易に使えるようにするためのライブラリです。

機能

構成

graph_netsは次のモジュールからなります。

graph_nets.graphs
グラフ(GraphsTuple)を定義
graph_nets.blocks
GNの要素（頂点・辺・全体）と集約の層を定義
graph_nets.modules
GNやその特殊なケースの層を定義
graph_nets.utils_tf
tensorflow用の次のような便利関数群を定義
graph_nets.utils_np
networkx形式とGraphsTuple形式とGraphsTuple風の辞書形式の相互変換の関数群等を定義

ユーザーがよく使いそうなのは、graph_nets.modulesやgraph_nets.utils_tf、graph_nets.utils_npでしょう。

graph_nets.modules

graph_nets.modulesはGN等の層を定義します。その使い方について擬似的なコードがREADMEのUsage exampleの箇所に書かれています。一部コメントを書き換えて引用します。

import graph_nets as gn
import sonnet as snt

# GNの入力の層を定義（この部分は擬似的に書かれています）
input_graphs = get_graphs()

# GNを構成。引数に頂点・辺・全体での処理内容（上の方の図のφ*の部分）や集約方法を与えます。
# この例ではそれぞれで多層パーセプトロン（全結合を数回のする奴）を行うものが指定されます。訓練時にこれらが訓練されます。
graph_net_module = gn.modules.GraphNetwork(
    edge_model_fn=lambda: snt.nets.MLP([32, 32]),
    node_model_fn=lambda: snt.nets.MLP([32, 32]),
    global_model_fn=lambda: snt.nets.MLP([32, 32]))

# Sonnetのモデルはこのようにmodule(inputs)のようにして使います
output_graphs = graph_net_module(input_graphs)

他にもRelationNetworkやSelfAttentionといった層があります。

公式サンプルコード

サンプルのコード・notebookがここにあります。バネ・ソート・最短経路を扱った例が並んでいます。また、そこで利用されるモデルの定義はmodels.pyにあり、使い方を知るのに役に立つと思います。

おまけ

Sonnet

GraphNetsはSonnetをベースに作成されたライブラリです。Sonnet自体もDeepMindのライブラリで、TensorFlowを少し便利に使う薄いラッパーです。kerasと異なり便利すぎる機能はありませんが、TensorFlowの世界からあまり離れること無くコードが書け、過剰な便利機能が無く、個人的には気に入っています。
Sonnetの名前を出すと、「生きていたのか」とかいう反応をされるのですが、もっと評価されていいライブラリと思います。

Sonnetを利用するのに覚えておくべき機能は、sonnet.AbstractModuleだけです。これだけ覚えれば、後はそれを使っている色々なパーツを単に使うだけです。次のように使います。

import sonnet as snt


class HogeLayer(snt.AbstractModule):
  def __init__(self, name=None):
    super().__init__(name=name)  # python2勢はよしなに
    with self._enter_variable_scope(): # スコープをいい感じにしてくれる呪文
      self.__conv = snt.Conv2D(128, 3)  # チャンネル数128,カーネルサイズ3の畳み込み

  def _build(self, inputs):
    return self._conv(inputs)

...
hoge = HogeLayer()  # ネットワークを設定
hoge_output0 = hoge(inputs0)  # ネットワークを構築。ここで"_build"が呼ばれています
hoge_output1 = hoge(inputs1)  # hoge.__convの畳み込みのパラメータが簡単に再利用できます

Sonnetが気楽なのは、_buildメソッドの返り値がそのまま出てくる点です。返り値がtf.Tensorになるようにすれば、通常のTensorflowの関数のように使えます。ネットワークの設定と構築の場所が分かれていることに注意すると容易に使えると思います。