グラフ向け深層学習ライブラリDeep Graph Library (DGL)の初歩の初歩

グラフ向けの深層学習ライブラリDeep Graph Library(DGL)の基本的な使い方について紹介します。[公式ドキュメント](https://docs.dgl.ai/)に事例やAPIの説明が詳細に載っていたり[チュートリアル](https://docs.dgl.ai/tutorials/models/index.html)も豊富にありますが、DGLの一番基本的な動作（だと個人的に思っている）ノードの特徴量のmessageとreduceという2つの処理について、丁寧に説明している記事がなかったので説明してみます。

## Deep Graph Library (DGL)とは？
New York UniversityとAWSが開発しているPytorch-basedの（？）グラフと対象としたDeep Learningのライブラリです。
画像や言語など従来よく研究されているデータ構造ではTensorFlow, Pytorch, Chainerなど有名なライブラリがあり、CNNやRNNなどが1つの関数（[公式](https://www.dgl.ai/pages/about.html)ではbuilding-blocksと言っている）になっていて、それを組み合わせることでモデルを作ることができます。しかし、グラフデータを対象としたライブラリではまだ決定版と言えるものがなく、DGLはその有力候補だと思います。

## 基本的な考え方
DGLの基本的な流れは以下の通りです。
1. グラフの各ノードに特徴量を割り当てる
2. あるノードから隣接ノードに1で割り当てた特徴量を送る（send）
3. 2で隣接ノードから送られてきた特徴量を集約する（recv）

ここの2,3が公式ドキュメントやチュートリアルで丁寧な説明が少し足りないかなと思いました。唯一書いてあるのは[Pagerankのチュートリアル](https://docs.dgl.ai/tutorials/basics/3_pagerank.html)の真ん中くらいにある図くらいだと思います。以下、基本的な流れの詳細を説明します。説明では[DGL at a Glance](https://docs.dgl.ai/tutorials/basics/1_first.html)という一番最初のチュートリアルをベースに、さらに簡単にしたグラフ構造で説明します。コードもこのチュートリアルをベースにしています。

## 事前準備
公式チュートリアルはグラフ分野で有名なKarateデータでやっていますが、もっと簡単に5つのノードで構成されるグラフを作ります。

```Python
import dgl

def build_sample_graph():
    g = dgl.DGLGraph()
    g.add_nodes(5)
    edge_list = [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (2, 3)]
    src, dst = tuple(zip(*edge_list))
    # 無向グラフにするためsrcとdstを入れ替えて両方の向きにエッジを設定
    g.add_edges(src, dst)
    g.add_edges(dst, src)

    return g
G = build_sample_graph()
```

可視化するとこんな感じです。

![graph_overview.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/37387/ac730b94-5ec0-b4e0-6717-eedab1fd6781.png)


## グラフの各ノードに特徴量を割り当てる
ここでは簡単に単位行列を各ノードの特徴量にします。ndataはノードの特徴量で、ここでは名前をhにしています。

```Python
G.ndata['h'] = torch.eye(5)
```
図にするとこんな感じです。
![graph_step1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/37387/18643bc9-880a-eda1-52cd-c9cf0b7908f8.png)


## 隣接ノードに特徴量を送る（send）
ソースとなるノードの特徴量（'h'）を隣接ノード（送り先では'msg'という名前で受信される）に送るコードは以下のように書きます。

```Python
def gcn_message(edges):
    return {'msg' : edges.src['h']}

G.send(G.edges(), gcn_message)
```

この時点では、グラフ（特徴量）は以下のようになっています。図の黒字は元々の特徴量、赤字はノード0から隣接ノードに送った特徴量、同様に黄色はノード1が送ったもの、緑字はノード2が送ったもの、青字はノード3が送ったもの、紫字はノード4が送ったものです。

![graph_step2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/37387/eb668475-22ed-5c89-2eec-59d0edad2968.png)


## 隣接ノードから送られてきた特徴量を集約する（recv）
最後に以下のようにして隣接ノードから送られてきた特徴量を集約します。送られてきた特徴量はnodeのmailboxというpropertyにあります。message処理で名前を'msg'としていたのでnodes.mailbox['msg']でデータにアクセスできます。
ここでは隣接ノードの特徴量の総和を新たにそのノードの特徴量（'h'）にセットしています。

```Python
def gcn_reduce(nodes):
    return {'h' : torch.sum(nodes.mailbox['msg'], dim=1)}

G.recv(G.nodes(), gcn_reduce)
```

## 重みの学習
重みの学習はPytorchの他のモデルと同様に以下のようにやればできます。

```Python
optimizer.zero_grad()
loss.backward()
optimizer.step()
```

ここまでの全体のコードはこちらにあげました。
https://gist.github.com/k1ochiai/cd0279ca79dd74e91a2b5e1187928adb


## 今後の発展
[Graph Convolutional Network](https://arxiv.org/abs/1609.02907)や[GraphSAGE](http://snap.stanford.edu/graphsage/)など最新のモデルを実装するには、この記事で説明したgcn_message, gcn_reduceという2つの関数を各手法になるように実装すればできます。

ちなみに[DGLのリポジトリ](https://github.com/dmlc/dgl/tree/master/examples/pytorch)に結構最新のモデルも実装されているので使えそうです。例えばICLR 2019で発表された「How Powerful are Graph Neural Networks?」も実装があります（[リンク](https://github.com/dmlc/dgl/tree/master/examples/pytorch/gin)）。