Graph Neural Networks で使われている計算の最小動作例
グラフデータを深層学習で扱う場合，Graph Neural Networks (GNNs) を利用することがあります[1]．これらを実装して何らかのアプリに利用するには pytorch-geometric や deep graph library (dgl) などを利用します．ライブラリを用いることで，通常の深層学習フレームワークを利用するのと同じようにニューラルネットワークを構成し，学習のフレームワーク (PyTorchならOptimizer，DataLoaderなど) を用いてパラメータの学習が可能です．
現代では，様々なフレームワークを利用してネットワークが簡単に書けるようになりました．一方でこれらのパッケージの実装を見ていると，数式との対応関係をぱっと理解するのがなかなか難しく「本当にこの式の計算がライブラリのこのLayerで計算されているのか？？？」というように，式と実装のギャップを埋めることが難しい場合もあります．そこで，簡単なグラフに対する例と数式を対応付けた最小の動作例 (Minimal working example) を書き記すことを目的に，この記事を書きました．
対象とする式
Kipf & Welling ICLR2017 など多くのGNN論文に出てくる式として，次式で表現される 1-layer GCNを対象とします．
ここで
- グラフは
G=(V, E), |V|=n - グラフ
G A\in\{0, 1\}^{n\times n} - 頂点
i\in V d \mathbf{x}_i\in\mathbb{R}^d -
n X\in\mathbb{R}^{n\times d} - モデル
f(A, X; W) W A X -
\hat{A} = A + I \hat{D}_{ii} = \sum_{j=1}^{n} \hat{A}_{ij} \hat{D} - 例えば
L H^{(0)}=X H{(l+1)} = f(A, H^{(l)}; W)
-
-
\sigma(\cdot) - たとえば
\mathrm{ReLU}(\cdot) \mathrm{softmax}(\cdot) - 今回は手抜きをしたので
\sigma(\cdot)
- たとえば
実装と観察
notebookはここにあります
例題
6頂点のグラフを考えます
V=\{0,1,2,3,4,5\} -
E=\{01, 02, 12, 23, 24, 45\} v_1v_2=\{v_1,v_2\} - 隣接行列
A=\begin{pmatrix}0&1&1&0&0&0\\1&0&1&0&0&0\\1&1&0&1&1&0\\0&0&1&0&0&0\\0&0&1&0&0&1\\0&0&0&0&1&0\end{pmatrix}
- 隣接行列
こちらのグラフに対して式を計算した場合と，dglを利用した場合の結果を比較して，動作を確認してみます．
numpyを利用する計算
次元を確認してみます．入力として各ノードに与えれる特徴ベクトルを集めた行列が
例
分かりやすさのために入力を固定しておきます．
-
X X=\begin{pmatrix}0&0&0\\0&0&1\\0&1&0\\0&1&1\\1&0&0\\1&1&0\end{pmatrix} - 重み
W 1 W=\begin{pmatrix}1&1\\1&1\\1&1\end{pmatrix}
後はnumpyで普通に行列計算をしましょう．ここでは後で使うために
n = 6
output_dim = 2
E = [(0, 1), (0, 2), (1, 2), (2, 3), (2, 4), (4, 5)]
A = np.zeros((n, n))
Ahat = A + np.eye(n)
Dhat = lambda x: np.diag(np.power(np.sum(Ahat, axis=1), x))
# 入力X
X = np.array([
[0, 0, 0], # 0
[0, 0, 1], # 1
[0, 1, 0], # 2
[0, 1, 1], # 3
[1, 0, 0], # 4
[1, 1, 0], # 5
])
d = X.shape[1]
W = np.ones((d, output_dim))
# 行列計算
Omat = Dhat(-1/2) @ Ahat @ Dhat(-1/2) @ X @ W
print(Omat)
"""出力
[[0.59153222 0.59153222]
[0.59153222 0.59153222]
[1.34885331 1.34885331]
[1.31622777 1.31622777]
[1.4080288 1.4080288 ]
[1.40824829 1.40824829]]
"""
まったく役に立たない計算ですが，
Deep Graph Library (dgl) を利用した例
同じ計算を dgl.nn.GraphConv を用いて計算してみます．モジュールの説明はここにあります．
グラフの作成
まずはdglでグラフを作成する必要があります．チュートリアル (A Blitz Introduction) がこちらにあります．チュートリアルでは基本的に有向グラフとしてグラフを扱い，無向グラフにするには辺を両方加えなさいと言われている気がするので，加えます．
dgl が他の形式からのグラフ構築をサポートしているので (このあたりに書かれています)，今回はnetworkxのオブジェクトから作成します．おまじないとして add_self_loop() を利用しています．
# networkxのオブジェクト
n = 6
nxg = nx.Graph()
nxg.add_nodes_from(range(n))
E = [(0, 1), (0, 2), (1, 2), (2, 3), (2, 4), (4, 5)]
for (u, v) in E:
nxg.add_edge(u, v)
# dgl用のグラフ
g = dgl.from_networkx(nxg)
g = dgl.add_self_loop(g) # おまじない (????????)
入力特徴の作成
次に
X = torch.tensor([[0, 0, 0], [0, 0, 1], [0, 1, 0],
[0, 1, 1], [1, 0, 0], [1, 1, 0]], dtype=torch.float32)
Graph Convolution Layerの作成
次はGCNLayerを作ります．これまで議論してきた式ではバイアス項は考えていないので，bias=False とします．先程まで出力の次元数を output_dim=2 としてきたので，そのように設定しています．
output_dim = 2
gcn = GraphConv(X.shape[1], output_dim, bias=False)
print(gcn)
# => GraphConv(in=3, out=2, normalization=none, activation=None)
レイヤーを作成したので，通常のLinear LayerなどをPyTorchで利用する雰囲気で
print(gcn(g, X))
"""出力
tensor([[ 0.2493, -0.2701],
[ 0.4356, -0.3569],
[ 0.3621, 0.9052],
[ 0.4356, -0.3569],
[ 1.1703, 0.3750],
[ 0.2992, 1.0887]], grad_fn=<GSpMMBackward>)
"""
パラメータの確認と固定
まずはOptimizerにパラメータを渡す際の雰囲気を思い出して，パラメータを確認してみましょう．
print(list(gcn.parameters()))
"""出力
[Parameter containing:
tensor([[ 0.3040, 0.7081],
[ 0.8723, -0.3016],
[ 0.7204, 1.0608]], requires_grad=True)]
"""
バイアス項を考えていないため，パラメータは
ここで強引にparametersを固定してみます．真面目にパラメータ書き換えを行う方法が色々とあると思いますが，今回のレイヤは gcn.weight が torch.nn.parameter.Parameter の情報を持っているので，直接いじって設定することにします．
gcn.weight = torch.nn.parameter.Parameter(torch.ones_like(org_weight))
print(gcn.weight)
"""出力
Parameter containing:
tensor([[1., 1.],
[1., 1.],
[1., 1.]], requires_grad=True)
"""
計算の再確認
ここまで来たのでもう一度実行例を確認してみます．
print(gcn(g, X))
"""出力
tensor([[0.5915, 0.5915],
[0.5915, 0.5915],
[1.3489, 1.3489],
[1.3162, 1.3162],
[1.4080, 1.4080],
[1.4082, 1.4082]], grad_fn=<MulBackward0>)
"""
同じ出力が得られました！これで安心して自分で行列演算を書くことなく，dgl.nn.GraphConv を使うことができそうです．
オプション引数 norm について
GraphConvレイヤーにはオプション引数 norm が none | right | both の3つから選択することができます．この結果をラムダ式で定義したnumpyの計算を使って確認してみましょう．
まずはdglを使った実装の出力です．初期値は both で，bothは「論文と同じ計算」となっているため，上に実装したnumpy実装の両方の累乗数を
# 3つレイヤーを作成してパラメータを固定する
gcn0 = GraphConv(X.shape[1], output_dim, bias=False, norm="none")
gcn1 = GraphConv(X.shape[1], output_dim, bias=False, norm="right")
gcn2 = GraphConv(X.shape[1], output_dim, bias=False, norm="both")
for gcn in [gcn0, gcn1, gcn2]:
gcn.weight = torch.nn.parameter.Parameter(torch.ones_like(org_weight))
# 出力する
print(gcn0(g, X))
print(gcn1(g, X))
print(gcn2(g, X))
"""出力
tensor([[2., 2.],
[2., 2.],
[5., 5.],
[3., 3.],
[4., 4.],
[3., 3.]], grad_fn=<GSpMMBackward>)
tensor([[0.6667, 0.6667],
[0.6667, 0.6667],
[1.0000, 1.0000],
[1.5000, 1.5000],
[1.3333, 1.3333],
[1.5000, 1.5000]], grad_fn=<MulBackward0>)
tensor([[0.5915, 0.5915],
[0.5915, 0.5915],
[1.3489, 1.3489],
[1.3162, 1.3162],
[1.4080, 1.4080],
[1.4082, 1.4082]], grad_fn=<MulBackward0>)
"""
次にnumpyの実装を用いてnoneとrightに相当する計算を再現してみます．実は none は両方に
# noneの再現
mat0 = (Dhat(0)) @ Ahat @ (Dhat(0)) @ X @ W
print(mat0)
"""出力
[[2. 2.]
[2. 2.]
[5. 5.]
[3. 3.]
[4. 4.]
[3. 3.]]
"""
一方で right ですが，説明によれば「If is right, divide the aggregated messages by each node's in-degrees, which is equivalent to averaging the received messages」と書いてあり，dglでの実装も別式で書かれていました．実装を解読するに，まずは左側の
-
GNNのお気持ちが気になる場合は例えばKips氏の https://tkipf.github.io/graph-convolutional-networks/ やPFN石黒さんの https://www.slideshare.net/pfi/20201023naistpfnishigurognnintroduction を見ましょう ↩︎
Discussion