📈

Graph Neural Networks で使われている計算の最小動作例

8 min read

グラフデータを深層学習で扱う場合,Graph Neural Networks (GNNs) を利用することがあります[1].これらを実装して何らかのアプリに利用するには pytorch-geometricdeep graph library (dgl) などを利用します.ライブラリを用いることで,通常の深層学習フレームワークを利用するのと同じようにニューラルネットワークを構成し,学習のフレームワーク (PyTorchならOptimizer,DataLoaderなど) を用いてパラメータの学習が可能です.

現代では,様々なフレームワークを利用してネットワークが簡単に書けるようになりました.一方でこれらのパッケージの実装を見ていると,数式との対応関係をぱっと理解するのがなかなか難しく「本当にこの式の計算がライブラリのこのLayerで計算されているのか???」というように,式と実装のギャップを埋めることが難しい場合もあります.そこで,簡単なグラフに対する例と数式を対応付けた最小の動作例 (Minimal working example) を書き記すことを目的に,この記事を書きました.

対象とする式

Kipf & Welling ICLR2017 など多くのGNN論文に出てくる式として,次式で表現される 1-layer GCNを対象とします.

f(A, X; W) = \sigma\left(\hat{D}^{-1/2} \hat{A} \hat{D}^{-1/2} X W\right)

ここで

  • グラフは G=(V, E), |V|=n です
  • グラフ G の隣接行列表現を A\in\{0, 1\}^{n\times n} とします
  • 頂点 i\in Vd 次元の特徴ベクトル \mathbf{x}_i\in\mathbb{R}^d を持ちます
  • n 頂点の特徴ベクトルをまとめた行列を X\in\mathbb{R}^{n\times d} で表します
  • モデル f(A, X; W) はパラメータ W で特徴付けられる計算式で,AX から新しい特徴行列を計算します
    • \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はここにあります

https://github.com/cocomoff/example_gcn

例題

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を利用する計算

次元を確認してみます.入力として各ノードに与えれる特徴ベクトルを集めた行列が (n, d) 次元です.グラフ計算の項は (n, n) 行列なので,右から(n, d) 行列がやってきて,更に右から重み行列 W(d, l) の形でやってくるとすると,最終的に (n, l) 次元の行列が出力されます.

分かりやすさのために入力を固定しておきます.

  • Xにノード番号の2進数表記で作った3次元の特徴ベクトル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で普通に行列計算をしましょう.ここでは後で使うために\hat{D}をラムダ式にしています.

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

まったく役に立たない計算ですが,f(A, X; W) を計算し,(n, l) 型の行列を取得することができました.

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を構築します.今,デフォルトでdglのbackendがpytorchになっているはずなので,pytorchの(n, d)型のテンソルにします.計算の都合上データの方は float32 を選択します (手元の環境だと指定しないと怒られました).

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で利用する雰囲気でXをforwardします.

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

(n, l)=(6, 2)型の出力が出てきているみたいですが,今回作成したGCNレイヤーは重みを内側に持っていてnumpyの計算みたいに固定されていないため,同じ出力にはなっていません (当たり前ですが…).

パラメータの確認と固定

まずはOptimizerにパラメータを渡す際の雰囲気を思い出して,パラメータを確認してみましょう.

print(list(gcn.parameters()))

"""出力
[Parameter containing:
tensor([[ 0.3040,  0.7081],
        [ 0.8723, -0.3016],
        [ 0.7204,  1.0608]], requires_grad=True)]
"""

バイアス項を考えていないため,パラメータはW\in\mathbb{R}^{d\times l}=\mathbb{R}^{3\times 2}です.大まかな計算の構成自体は,だいたいnumpyのものと似たようになっていそうで,問題なさそう(?)ですね.

ここで強引にparametersを固定してみます.真面目にパラメータ書き換えを行う方法が色々とあると思いますが,今回のレイヤは gcn.weighttorch.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実装の両方の累乗数を -1/2 に設定したものになります.

# 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 は両方に 0 を利用する場合に相当します.

# 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での実装も別式で書かれていました.実装を解読するに,まずは左側の\hat{D}は不要になり,右側の\hat{D}の部分だけをindegreeのベクトルで割ればいいはずなのですが,あまり正確に分かりませんでした (?).

脚注
  1. GNNのお気持ちが気になる場合は例えばKips氏の https://tkipf.github.io/graph-convolutional-networks/ やPFN石黒さんの https://www.slideshare.net/pfi/20201023naistpfnishigurognnintroduction を見ましょう ↩︎

Discussion

ログインするとコメントできます