用 Pytorch 实现一个简单的图卷积神经网络

深度学习在序列化数据的处理上取得了良好的效果, 但对于非欧空间中的数据处理问题却长时间难以突破. 在 2017 年, Kipf 等人在图神经网络上利用信号处理中的谱方法成功实现了卷积操作, 将卷积神经网络成功应用在图上, 即图卷积神经网络 (Graph Convolutional Network) 开启了在图上进行深度学习的新时代.

本文通过利用 GCN 实现一个简单的图数据集的分类问题展示图神经网络的效果

本文不会涉及到任何 GCN 的理论知识, 仅适用于想体验 GCN 的效果的朋友们尝鲜.

数据集介绍

Zachary karate club 是一个美国大学空手道俱乐部的社会网络图, 网络中包含 $34$34 个节点和 $78$78 条边, 分别表示 $34$34 位成员和他们之间的关系, 若两成员间有边存在, 则该两位成员存在友谊关系. 详细信息可以在这里看到. 在该网络中, 俱乐部的成员最终分化为由管理员和教练带头的两个组织.

本文使用 networkx 导入数据集.

import networkx as nx
G = nx.karate_club_graph()
print(G.number_of_nodes()) # 34
print(G.number_of_edges()) # 78

GCN 介绍

与普通的卷积神经网络一样, GCN 的原理也是利用堆叠一定卷积层来实现, 层与层之间同样有**函数.

首先介绍拉普拉斯矩阵 (Laplacian matrix), 定义为
$L = D - A$L=D−A
其中, $A$A 为图的邻阶矩阵, $D$D 为图的度矩阵, $D$D 的定义为
KaTeX parse error: No such environment: equation at position 8: \begin{̲e̲q̲u̲a̲t̲i̲o̲n̲}̲ \begin{aligned…
本文使用的是泛化的拉普拉斯矩阵(Generalized Laplacian matrix)
$L_g = D^{-1/2}LD^{-1/2} = I - D^{-1/2}AD^{-1/2}$Lg​=D−1/2LD−1/2=I−D−1/2AD−1/2
其中, $I$I 为单位阵.

为了增加网络学习时的稳定性, 作者对 $L_g$Lg​ 做了归一化处理, 即
$L_g = L_g + I$Lg​=Lg​+I
则图卷积层可以表示为
$X' = \sigma(L_gXW)$X′=σ(Lg​XW)
$X$X 为图上结点的特征 (feature) 值, 由于我们使用的数据集没有特征信息, 所以使用 ont-hot 代替. $W$W 是用于线性变换的参数.

算法实现

首先导入需要的一些依赖项

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import torch.nn.init as init

拉普拉斯的归一化操作

def norm(adj):
    adj += np.eye(adj.shape[0]) # 为每个结点增加自环
    degree = np.array(adj.sum(1)) # 为每个结点计算度
    degree = np.diag(np.power(degree, -0.5))
    return degree.dot(adj).dot(degree)

卷积层的设计

class GraphConvolution(nn.Module):
    def __init__(self, input_size, output_size):
        super(GraphConvolution, self).__init__()
        self.linear = nn.Linear(input_size, output_size)

    def forward(self, adj, features):
        out = torch.mm(adj, features)
        out = self.linear(out)
        return out

GCN 设计

class GCN(nn.Module):
    def __init__(self, input_size=34, hidden_size=5):
        super(GCN, self).__init__()
        self.gcn1 = GraphConvolution(input_dim, hidden_size)
        self.gcn2 = GraphConvolution(hidden_size, 2)
    
    def forward(self, adj, features):
        out = F.relu(self.gcn1(adj, features))
        out = self.gcn2(adj, out)
        return out

训练参数

LEARNING_RATE = 0.1
WEIGHT_DACAY = 5e-4
EPOCHS = 50
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

训练数据准备

features = np.eye(34, dtype=np.float)

y = np.zeros(G.number_of_nodes())
for i in range(G.number_of_nodes()):
    if G.nodes[i]['club'] == 'Mr. Hi':
        y[i] = 0
    else:
        y[i] = 1
        
adj = np.zeros((34, 34)) # 邻阶矩阵
for k, v in G.adj.items():
    for item in v.keys():
        adj[k][item] = 1
adj = norm(adj)

分别将其转换为 tensor

features = torch.tensor(features, dtype=torch.float).to(DEVICE)
y = torch.tensor(y, dtype=torch.long).to(DEVICE)
adj = torch.tensor(adj, dtype=torch.float).to(DEVICE)

开始训练

为了展示 GCN 强大的能力, 整个训练过程只使用管理员和教练两个结点进行.

def train():
    for epoch in range(EPOCHS):
        out = net(adj, features)
        mask = [False if x != 0 and x != 33 else True for x in range(34)] # 只选择管理员和教练进行训练
        l = loss(out[mask], y[mask])
        optimizer.zero_grad()
        l.backward()
        optimizer.step()

        print(f"epoch: {epoch}, loss: {l.item()}")

训练结束后顺便可视化一下

r = net(adj, features).cpu().detach().numpy()
import matplotlib.pyplot as plt
fig = plt.figure()
for i in range(34):
    plt.scatter(r[i][0], r[i][1], color="r" if y[i] == 0 else 'b')
plt.show()