图神经网络——超大图上的节点表征学习

图卷积网络(GCN)已经成功地应用于许多基于图形的应用，然而大规模的GCN的训练仍然具有挑战性。目前基于SGD的算法要么面临着随GCN层数呈指数增长的高计算成本，要么面临着保存整个图形和每个节点的embedding到内存的巨大空间(显存)需求。于是论文Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Network提出了Cluster-GCN方法来解决超大图的训练问题。

1.Cluster-GCN方法简单概括

利用图节点聚类算法将一个图的节点划分为$c$个簇，每一次选择几个簇的节点和这些节点对应的边构成一个子图，然后对子图做训练。
由于是利用图节点聚类算法将节点划分为多个簇，所以簇内边的数量要比簇间边的数量多得多，所以可以提高表征利用率，并提高图神经网络的训练效率。
每一次随机选择多个簇来组成一个batch，这样不会丢失簇间的边，同时也不会有batch内类别分布偏差过大的问题。
基于小图进行训练，不会消耗很多内存空间，于是我们可以训练更深的神经网络，进而可以达到更高的精度。

2.GCN回顾

论文2017. Semi-supervised classification with graph convolutional networks中提出了GCN。作者定义了一种类似卷积的运算，称为谱图卷积。这使得CNN可以直接在图形上操作。GCN中的每一层通过考虑相邻节点的embedding，来更新Graph中的每个节点的特征向量表示。GCN的逐层正向传播可以总结为：
$$
X^{\left( l+1 \right)}=f\left( X^l,A \right) =\sigma \left( \tilde{D}^{-\frac{1}{2}}\tilde{A}D^{-\frac{1}{2}}X^{\left( l \right)}W^{\left( l \right)} \right)
$$

$X$是所有节点的特征向量构成的特征矩阵。
$X^l$和$X^{l+1}$分别是$l$层的输入和输出矩阵，$X^l$也是第$l$层对于所有节点的embedding。
A是图的邻接矩阵。
${A}=A+I_N$是带有自环的无向图的邻接矩阵。
$I_N$是单位矩阵。
$\tilde{D}{ii}=\sum{\tilde{A}{ij}}$是带有自环的无向图的度矩阵，是一个对角矩阵。
$W^{(l)}$是一个可训练权重矩阵或参数矩阵。
$\sigma(\cdot)$是激活函数，通常被设为ReLU。

为了简化计算，假设所有层的表征维度都是F，上式写为：
$$
Z^{(l+1)}=A^{\prime} X^{(l)} W^{(l)}, X^{(l+1)}=\sigma\left(Z^{(l+1)}\right)
$$

$A^{\prime}$是归一化和规范化后的邻接矩阵

将GCN用于半监督节点分类时，目标是通过最小化损失函数来学习上式中的权重矩阵：
$$
\mathcal{L}=\frac{1}{\left|\mathcal{Y}{L}\right|} \sum{i \in \mathcal{Y}{L}} \operatorname{loss}\left(y{i}, z_{i}^{L}\right)
$$

$\mathcal{Y}_{\mathcal{L}}$是节点类别。
$z_{i}^{(L)}$是$Z^{(L)}$的第$i$行，表示节点$i$的最终层预测。
$y_{i}$为节点$i$的真实类别。

3.Cluster-GCN的速度
原始的GCN的训练使用的是全批量梯度下降(Full Gradient Descent)，它的计算和内存成本很高:
在内存方面，通过反向传播来计算full gradient需要存储所有的embedding矩阵，这需要O(NFL)的内存空间（F是特征数量，N是结点数量，L是网络层数）。
在收敛速度方面，由于在每个epoch模型才更新一次，所以需要很多个epoch才会使模型达到收敛。

mini-batch SGD可以提高GCN的训练速度，其不需要计算整个梯度，只需要计算每次更新的mini-batch的梯度。使用大小为$b=|\mathcal{B}|$的$\mathcal{B} \subseteq[N]$来表示一个batch的节点索引。SGD的每一步都将计算梯度估计值$\frac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} \nabla \operatorname{loss}\left(y_{i}, z_{i}^{(L)}\right)$来进行参数更新。

尽管在每个epoch收敛得更快，但SGD在GCN训练上引入了另一个计算开销（邻域扩展），这使得SGD的每个epoch时间内比全梯度下降（full gradient descent）要慢得多——计算一个节点$i$的梯度，需要节点$i$的embedding，而节点$i$的embedding的依赖于前面的层里的它的邻居的embeddings，为了提取前面的层中它的邻居节点embeddings，还需要进一步聚合每个邻居节点的邻居节点的embeddings。假设一个图神经网络有$L+1$层，节点的平均的度为$d$。为了得到节点$i$的梯度，需要聚合图上$O\left(d^{L}\right)$的节点的表征，与权重矩阵$W^{(l)}$相乘，所以计算任意节点表征的时间开销是$O\left(F^{2}\right)$。所以SGD一个节点的梯度的计算需要$O\left(d^{L} F^{2}\right)$的时间。

4.Cluster-GCN详解

在mini-batch SGD更新中，设计一个batch和相应的计算sub graph来最大限度地提高embedding utilization。Cluster-GCN通过将embedding utilization连接到一个聚类目标上来实现。

对于一个Graph$G$，我们将其节点划分为$c$个簇：$\mathcal{V}=\left[\mathcal{V}{1}, \cdots \mathcal{V}{c}\right]$，其中$\mathcal{V}{t}$由第$t$个簇中的节点组成，得到$c$个子图:
$$
\bar{G}=\left[G{1}, \cdots, G_{c}\right]=\left[\left{\mathcal{V}{1}, \mathcal{E}{1}\right}, \cdots,\left{\mathcal{V}{c}, \mathcal{E}{c}\right}\right]
\notag
$$
其中$\mathcal{E}{t}$只由$\mathcal{V}{t}$中的节点之间的边组成。经过节点重组，邻接矩阵被划分为大小为$c^{2}$的块矩阵，如下所示
$$
A=\bar{A}+\Delta=\left[\begin{array}{ccc}
A_{11} & \cdots & A_{1 c} \
\vdots & \ddots & \vdots \
A_{c 1} & \cdots & A_{c c}
\end{array}\right]
$$
其中
$$
\bar{A}=\left[\begin{array}{ccc}
A_{11} & \cdots & 0 \
\vdots & \ddots & \vdots \
0 & \cdots & A_{c c}
\end{array}\right], \Delta=\left[\begin{array}{ccc}
0 & \cdots & A_{1 c} \
\vdots & \ddots & \vdots \
A_{c 1} & \cdots & 0
\end{array}\right]
$$

对角线上的块$A_{t t}$是大小为$\left|\mathcal{V}{t}\right| \times\left|\mathcal{V}{t}\right|$的邻接矩阵，它由$G_{t}$内部的边构成。
$\bar{A}$是图$\bar{G}$的邻接矩阵。
$A_{s t}$由两个簇$\mathcal{V}{s}$和$\mathcal{V}{t}$之间的边构成。
$\Delta$是由$A$的所有非对角线块组成的矩阵。

用块对角线邻接矩阵$\bar{A}$去近似邻接矩阵$A$之后，完整的损失函数可以根据batch分解成多个部分之和：

Cluster-GCN使用了Graph聚类算法来划分Graph。Graph聚类的方法，如metis和graclus等，旨在在Graph中的顶点上构建分区，使得簇内连接远大于簇间连接，从而更好的捕捉聚类和区分结构。

5.Cluster-GCN实践

PyG库中提供了Cluster-GCN的接口，可以像训练普通神经网络一样在超大图上训练图神经网络。

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 #导入库和数据集 import torch import torch.nn.functional as F from torch.nn import ModuleList from torch_geometric.nn import SAGEConv from tqdm import tqdm from torch_geometric.datasets import Reddit, Reddit2 from torch_geometric.data import ClusterData, ClusterLoader, NeighborSampler dataset = Reddit('../dataset/Reddit') data = dataset[0] print(dataset.num_classes) print(data.num_nodes) print(data.num_edges) print(data.num_features)


`1 2 3 4 5 6 7 8 9 10 11 12 13 14 15`	`#导入库和数据集 import torch import torch.nn.functional as F from torch.nn import ModuleList from torch_geometric.nn import SAGEConv from tqdm import tqdm from torch_geometric.datasets import Reddit, Reddit2 from torch_geometric.data import ClusterData, ClusterLoader, NeighborSampler dataset = Reddit('../dataset/Reddit') data = dataset[0] print(dataset.num_classes) print(data.num_nodes) print(data.num_edges) print(data.num_features)`

结果：
41, 232965, 114615892, 602


`1 2 3 4 5 6`	`#图节点聚类并生成数据加载器 cluster_data = ClusterData(data, num_parts=1500, recursive=False, save_dir=dataset.processed_dir) train_loader = ClusterLoader(cluster_data, batch_size=20, shuffle=True, num_workers=12) subgraph_loader = NeighborSampler(data.edge_index, sizes=[-1], batch_size=1024, shuffle=False,num_workers=12)`

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 #构建网络 class Net(torch.nn.Module): def __init__(self, in_channels, out_channels): super(Net, self).__init__() self.convs = ModuleList( [SAGEConv(in_channels, 128), SAGEConv(128, out_channels)]) def forward(self, x, edge_index): for i, conv in enumerate(self.convs): x = conv(x, edge_index) if i != len(self.convs) - 1: x = F.relu(x) x = F.dropout(x, p=0.5, training=self.training) return F.log_softmax(x, dim=-1) def inference(self, x_all): pbar = tqdm(total=x_all.size(0) * len(self.convs)) pbar.set_description('Evaluating') for i, conv in enumerate(self.convs): xs = [] for batch_size, n_id, adj in subgraph_loader: edge_index, _, size = adj.to(device) x = x_all[n_id].to(device) x_target = x[:size[1]] x = conv((x, x_target), edge_index) if i != len(self.convs) - 1: x = F.relu(x) xs.append(x.cpu()) pbar.update(batch_size) x_all = torch.cat(xs, dim=0) pbar.close() return x_all


`1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38`	#构建网络 class Net(torch.nn.Module): def __init__(self, in_channels, out_channels): super(Net, self).__init__() self.convs = ModuleList( [SAGEConv(in_channels, 128), SAGEConv(128, out_channels)]) def forward(self, x, edge_index): for i, conv in enumerate(self.convs): x = conv(x, edge_index) if i != len(self.convs) - 1: x = F.relu(x) x = F.dropout(x, p=0.5, training=self.training) return F.log_softmax(x, dim=-1) def inference(self, x_all): pbar = tqdm(total=x_all.size(0) * len(self.convs)) pbar.set_description('Evaluating') for i, conv in enumerate(self.convs): xs = [] for batch_size, n_id, adj in subgraph_loader: edge_index, _, size = adj.to(device) x = x_all[n_id].to(device) x_target = x[:size[1]] x = conv((x, x_target), edge_index) if i != len(self.convs) - 1: x = F.relu(x) xs.append(x.cpu()) pbar.update(batch_size) x_all = torch.cat(xs, dim=0) pbar.close() return x_all

此神经网络的forward函数的定义与普通的图神经网络并无区别，而inference方法应用于推理阶段，为了获取更高的预测精度，所以使用subgraph_loader。

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 #训练、验证、测试 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = Net(dataset.num_features, dataset.num_classes).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.005) def train(): model.train() total_loss = total_nodes = 0 for batch in train_loader: batch = batch.to(device) optimizer.zero_grad() out = model(batch.x, batch.edge_index) loss = F.nll_loss(out[batch.train_mask], batch.y[batch.train_mask]) loss.backward() optimizer.step() nodes = batch.train_mask.sum().item() total_loss += loss.item() * nodes total_nodes += nodes return total_loss / total_nodes @torch.no_grad() def test(): # Inference should be performed on the full graph. model.eval() out = model.inference(data.x) y_pred = out.argmax(dim=-1) accs = [] for mask in [data.train_mask, data.val_mask, data.test_mask]: correct = y_pred[mask].eq(data.y[mask]).sum().item() accs.append(correct / mask.sum().item()) return accs for epoch in range(1, 31): loss = train() if epoch % 10 == 0: train_acc, val_acc, test_acc = test() print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, ' f'Val: {val_acc:.4f}, test: {test_acc:.4f}') else: print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}')


`1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44`	#训练、验证、测试 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = Net(dataset.num_features, dataset.num_classes).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.005) def train(): model.train() total_loss = total_nodes = 0 for batch in train_loader: batch = batch.to(device) optimizer.zero_grad() out = model(batch.x, batch.edge_index) loss = F.nll_loss(out[batch.train_mask], batch.y[batch.train_mask]) loss.backward() optimizer.step() nodes = batch.train_mask.sum().item() total_loss += loss.item() * nodes total_nodes += nodes return total_loss / total_nodes @torch.no_grad() def test(): # Inference should be performed on the full graph. model.eval() out = model.inference(data.x) y_pred = out.argmax(dim=-1) accs = [] for mask in [data.train_mask, data.val_mask, data.test_mask]: correct = y_pred[mask].eq(data.y[mask]).sum().item() accs.append(correct / mask.sum().item()) return accs for epoch in range(1, 31): loss = train() if epoch % 10 == 0: train_acc, val_acc, test_acc = test() print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, ' f'Val: {val_acc:.4f}, test: {test_acc:.4f}') else: print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}')

结果：
Epoch: 10, Loss: 0.2893, Train: 0.9620, Val: 0.9516, test: 0.9496
Epoch: 20, Loss: 0.2478, Train: 0.9664, Val: 0.9546, test: 0.9560
Epoch: 30, Loss: 0.2397, Train: 0.9689, Val: 0.9515, test: 0.9502

在训练过程中，使用train_loader获取batch，每次根据多个簇组成的batch进行训练；在验证阶段，使用subgraph_loader，在计算一个节点的表征时会计算该节点的距离从$0$到$L$的邻接节点，可以更好地测试神经网络的性能。

参考资料

1.datawhale-GNN开源学习资料
2.Cluster-GCN论文