超大图上的节点表征学习
引用自Datawhale
https://github.com/datawhalechina/team-learning-nlp/tree/master/GNN
前面文章:
什么是超大图?
- 节点以及边达到千万或上亿大数据规模,表征这些节点对内存,以及GPU计算带来了性能挑战。
- 图神经网络层数,如计算机视觉里面深层很多的网络,图神经网络层数过多带来的过度平滑问题。
本文知识主要引用论文Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Network
Cluster-GCN概括
为了解决普通的方法无法在超大图上做节点表征学习的问题,Cluster-GCN论文提出:
- 利用图节点聚类算法将一个图的节点划分为 c c c个簇,每一次选择几个组的节点和这些节点对应的边构成一个子图,然后对子图做训练。
- 由于是利用图节点聚类算法将节点划分为多个簇,所以簇内边的数量要比簇间边的数量多得多。
- 基于小图进行训练,不会消耗很多内存空间,于是我们可以训练更深的神经网络,从而得到更高的精度。
节点表征学习背景
给定一个图
G
=
(
V
,
E
,
A
)
G=(\mathcal{V}, \mathcal{E}, A)
G=(V,E,A),它由
N
=
∣
V
∣
N=|\mathcal{V}|
N=∣V∣个节点和
∣
E
∣
|\mathcal{E}|
∣E∣条边组成,其邻接矩阵记为
A
A
A。每一个节点都关联一个
F
F
F维的属性向量。
X
∈
R
N
×
F
X \in \mathbb{R}^{N \times F}
X∈RN×F表示
N
N
N的特征。一个
L
L
L层的图神经网络由
L
L
L个图卷积层组成,每一层都通过聚合上一层的图中节点的邻居的嵌入来构建每个节点的嵌入:
Z
(
l
+
1
)
=
A
′
X
(
l
)
W
(
l
)
,
X
(
l
+
1
)
=
σ
(
Z
(
l
+
1
)
)
(1)
Z^{(l+1)}=A^{\prime} X^{(l)} W^{(l)}, X^{(l+1)}=\sigma\left(Z^{(l+1)}\right) \tag{1}
Z(l+1)=A′X(l)W(l),X(l+1)=σ(Z(l+1))(1)
其中
X
(
l
)
∈
R
N
×
F
l
X^{(l)} \in \mathbb{R}^{N \times F_{l}}
X(l)∈RN×Fl是第
l
l
l层
N
N
N个节点的嵌入(embedding),并且有
X
(
0
)
=
X
X^{(0)}=X
X(0)=X。
A
′
A^{\prime}
A′是归一化和规范化后的邻接矩阵,
W
(
l
)
∈
R
F
l
×
F
l
+
1
W^{(l)} \in \mathbb{R}^{F_{l} \times F_{l+1}}
W(l)∈RFl×Fl+1是特征转换矩阵,也就是要学习的参数。为了简单起见,我们假设所有层的特征维度都是一样的,即
(
F
1
=
⋯
=
F
L
=
F
)
\left(F_{1}=\cdots=F_{L}=F\right)
(F1=⋯=FL=F)。激活函数
σ
(
⋅
)
\sigma(\cdot)
σ(⋅)通常被设定为ReLU
。
当图神经网络应用于半监督节点分类任务时,训练的目标是通过最小化损失函数来学习公式(1)中的权重矩阵:
L
=
1
∣
Y
L
∣
∑
i
∈
Y
L
loss
(
y
i
,
z
i
L
)
(2)
\mathcal{L}=\frac{1}{\left|\mathcal{Y}_{L}\right|} \sum_{i \in \mathcal{Y}_{L}} \operatorname{loss}\left(y_{i}, z_{i}^{L}\right) \tag{2}
L=∣YL∣1i∈YL∑loss(yi,ziL)(2)
其中,
y
L
y_{L}
yL包含所有被标记节点的标签;
z
i
(
L
)
z_{i}^{(L)}
zi(L)是
Z
(
L
)
Z^{(L)}
Z(L)的第
i
i
i行,表示节点
i
i
i的最终预测,并且其对应真实标签为
y
i
y_{i}
yi。
以往训练方法的瓶颈
以往的训练方法需要计算所有节点的嵌入以及所有(有标签)节点的损失产生的梯度,这种训练方式需要非常巨大的计算开销和GPU内存开销:在内存方面,通过反向传播计算(2)的全部损失的梯度需要存储所有的嵌入矩阵
{
Z
(
l
)
}
l
=
1
L
\left\{Z^{(l)}\right\}_{l=1}^{L}
{Z(l)}l=1L,这需要
O
(
N
F
L
)
O(N F L)
O(NFL)空间;在收敛速度方面,由于神经网络在每个epoch中只更新一次,所以训练需要更多的epoch才能达到收敛。
最近的一些工作证明采用mini-batch SGD的方式训练可以提高图神经网络的训练速度并减少GPU内存需求。在参数更新中,SGD不需要计算完整的梯度,而只需要基于mini-batch计算梯度。我们使用
B
⊆
[
N
]
\mathcal{B} \subseteq[N]
B⊆[N]来表示一个节点索引batch,其大小为
b
=
∣
B
∣
b=|\mathcal{B}|
b=∣B∣。SGD的每一步都将计算梯度估计值
1
∣
B
∣
∑
i
∈
B
∇
loss
(
y
i
,
z
i
(
L
)
)
\frac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} \nabla \operatorname{loss}\left(y_{i}, z_{i}^{(L)}\right)
∣B∣1∑i∈B∇loss(yi,zi(L))来进行参数更新。尽管在epoches数量相同的情况下,采用SGD方式进行训练,收敛速度可以更快,但此种训练方式会引入额外的时间开销,这使得该训练方式相比全梯度下降的训练方式每个epoch的时间开销要大得多。
Cluster-GCN实践
PyG为Cluster-GCN提出的训练方式和神经网络的构建提供了良好的支持。我们无需在意图节点是如何被划分成多个簇的,PyG提供的接口允许我们像训练普通神经网络一样在超大图上训练图神经网络。
数据集分析
from torch_geometric.datasets import Reddit
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
# 114615873
# 602
可以看到该数据集包含41个分类任务,232,965个节点,114,615,873条边,节点维度为602维。
图节点聚类与数据加载器生成
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)
train_loader
,此数据加载器遵循Cluster-GCN提出的方案,图节点首先被聚类,返回的一个batch由多个簇组成。
subgraph_loader
,使用此数据加载器不对图节点聚类,计算一个batch中的节点的嵌入需要计算该batch中的所有节点的距离从
0
0
0到
L
L
L的邻居节点。
构造图神经网络的规范
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')
# Compute representations of nodes layer by layer, using *all*
# available edges. This leads to faster computation in contrast to
# immediately computing the final representations of each batch.
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
两个方法。forward
函数的定义与普通的图神经网络并无区别。inference
方法应用于推理阶段,为了获取更高的预测精度,所有使用subgraph_loader
。
训练、验证与测试
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 % 5 == 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}')
可见在训练过程中,我们从train_loader
获取batch,每次只根据一个batch进行参数训练。但在测试节点,我们不做图节点聚类,因此不会有边丢失的情况。
在此篇文章中,我们学习了超大图上的节点表征学习面临的挑战,然后学习了应对这一挑战的Cluster-GCN提出的方法,在实践部分我们还学习了使用Cluster-GCN提出的方案进行超大图节点分类的方法,以及图神经网络的设计规范。
作业
- 尝试将数据集切分成不同数量的簇进行实验,然后观察结果并进行比较。
数据集由于电脑配置不高,只能用Cora数据集做实验。
50簇
100簇
150簇
300簇
500簇
可以看到随着簇数增加,准确率有轻微增加,而后下降,时间则是先减少,后增加。
参考资料
- 提出Cluster-GCN的论文:Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Network
- [9]: Thomas N. Kipf and Max Welling. 2017. Semi-Supervised Classification with Graph Convolutional Networks. In ICLR.