原教程参考：torch_geometric官方教程

Basic property of dataset

1
2
3
from torch_geometric.datasets import TUDataset

dataset = TUDataset(root='D:/data', name='ENZYMES', use_node_attr = True)

上图下载并读取名为”ENZYMES”的数据集，数据集中具体的属性大小与内容如下所示：

1
2
3
4
5
dataset
>>> ENZYMES(600)

dataset[0]
>>> Data(edge_index=[2, 168], x=[37, 21], y=[1])

“ENZYMES(600)”代表共有600个样本，edge_index代表节点与节点之间的连接关系，同列的两个节点相互连接，有向图代表共有168条边，如果为无向图，则共有168/2条边。x代表每个节点的属性，共有37个节点，每个节点有21个特征。y代表图的结果。

Mini-batches

PyG通过将多个图的邻接矩阵A置于对角线上形成一个稀疏的对角矩阵，同时将x的属性和y堆叠为矩阵来形成一个mini batch用以并行计算。
$$
\mathbf{A}=
\begin{bmatrix}
\mathbf{A}_{1} & & \\
& \ddots & \\
& & \mathbf{A}_{n}
\end{bmatrix},
\quad
\mathbf{X}=
\begin{bmatrix}
\mathbf{X}_{1} \\
\vdots \\
\mathbf{X}_{n}
\end{bmatrix},
\quad
\mathbf{Y}=
\begin{bmatrix}
\mathbf{Y}_{1} \\
\vdots \\
\mathbf{Y}_{n}
\end{bmatrix}
$$

1
2
3
4
5
6
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in loader:
    batch
    break
>>> DataBatch(edge_index=[2, 3856], x=[1102, 21], y=[32], batch=[1102], ptr=[33])

代表着共有1102个节点，edge_index中显示着这个大图中的连接关系，共有3856/2(无向图)条边。ptr中存放每个图的开始与结束的节点；batch是一个将每个节点映射到其对应图的向量。

Data Transform

这里我们并没有使用ENZYMES数据集，而是使用ShapeNet数据集，因为该数据在输入到模型内时必须使用一些策略来处理。

1
2
3
4
5
6
from torch_geometric.datasets import ShapeNet

dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'])

dataset[0]
>>> Data(pos=[2518, 3], y=[2518])

上图中ShapeNet为一个3D点图，每个节点具有一个3维坐标，为了能够将其转化为图格式，我们不妨让每个点与k个最近的点相连接。所以可以使用pre_transform参数。

1
2
3
4
5
6
7
8
9
10
11
import torch_geometric.transforms as T
from torch_geometric.datasets import ShapeNet

dataset = ShapeNet(
	root='/tmp/ShapeNet', 
	categories=['Airplane'], 
	pre_transform=T.KNNGraph(k=6)
)

dataset[0]
>>> Data(edge_index=[2, 15108], pos=[2518, 3], y=[2518])

同时可以使用transform参数来引入随机扰动。

1
2
3
4
5
6
dataset = ShapeNet(
	root='/tmp/ShapeNet', 
	categories=['Airplane'],
	pre_transform=T.KNNGraph(k=6),
	transform=T.RandomJitter(0.01)
)

GCN model

torch_geometric官方教程
最简单的模型就是GCN模型，我们选取该模型作为复现的模型。数据集选用Cora数据集。这里Cora数据集的存储方式不同于ENZYMES数据集，其将所有数据均存放于列表中的第一个元素中，并用train_mask, val_mask, test_mask分别表示训练集、验证集、测试集。所以在训练、测试过程中，均采用mask进行切片后再处理。

1
2
3
4
5
from torch_geometric.datasets import Planetoid
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = Planetoid(root='D:/data/Graph_data', name='Cora')
data = dataset[0].to(device)

GCN模型中使用的图卷积操作来自于Semi-supervised Classification with Graph Convolutional Networks，以下为算法介绍：
$$
\mathbf{X}^{\prime} = \mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}}
\mathbf{\hat{D}}^{-1/2} \mathbf{X} \mathbf{\Theta},
$$
where $\mathbf{\hat{A}} = \mathbf{A} + \mathbf{I}$ denotes the adjacency matrix with inserted self-loops and $\hat{D}_{ii} = \sum_{j=0} \hat{A}_{ij}$ its diagonal degree matrix. The adjacency matrix can include other values than $1$ representing edge weights via the optional $\text{edge_weight}$ tensor.

Its node-wise formulation is given by:
$$
\mathbf{x}^{\prime}_i = \mathbf{\Theta}^{\top} \sum_{j \in
\mathcal{N}(i) \cup { i }} \frac{e_{j,i}}{\sqrt{\hat{d}_j
\hat{d}_i}} \mathbf{x}_j
$$
with $\hat{d}_i = 1 + \sum_{j \in \mathcal{N}(i)} e_{j,i}$, where $e_{j,i}$ denotes the edge weight from source node j to target i (default 1.0)。
其中，$\Theta$为可学习的矩阵，代表着一次全链接层，用于将特征混合后的矩阵转化为特定维度的特征。该算法如下：
1.邻接矩阵归一化：$\mathbf{\hat{D}}^{-1/2}\mathbf{\hat{A}}\mathbf{\hat{D}}^{-1/2}$
2.特征混合：$\mathbf{X}$
3.特征变换：$\Theta$
最终得到的矩阵形状为$[\text{node_num}, \text{out_channel}]$

模型创建

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

import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 64)
        self.bn1 = torch.nn.BatchNorm1d(64)
        
        self.conv2 = GCNConv(64, 64)
        self.bn2 = torch.nn.BatchNorm1d(64)

        self.linear = torch.nn.Linear(64, dataset.num_classes)
        self.reset_parameters()
        
    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = self.conv1(x, edge_index)
        x = self.bn1(x)
        x = F.relu(x)
        x = F.dropout(x, p=0.1, training=self.training)

        x = self.conv2(x, edge_index)
        x = self.bn2(x)
        x = F.relu(x)
        x = F.dropout(x, p=0.1, training=self.training)
		        
		# x = global_mean_pool(x, batch)
        # x = self.linear(x)
        
        return F.log_softmax(x, dim=1)

model = GCN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, weight_decay=5e-4)
model.train()

上图模型为针对节点分类的模型，所以最后x经过两层图卷积层及批归一化和激活函数。如果对于图分类任务，则仅需将注释的代码加入，用于进行所有节点的信息聚合与映射至目标空间。

模型训练

1
2
3
4
5
6
7
8
9
10
11
12
13
acc_list = []
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
        
    model.eval()
    pred = model(data).argmax(dim=1)
    correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
    acc = int(correct) / int(data.test_mask.sum())
    print(f'Accuracy: {acc:.4f}')

Cover image icon by Dewi Sari from Flaticon