GNN图神经网络

图神经网络(Graph Neural Networks, GNNs)是一种专门处理图结构数据的深度学习模型。它可以很好地捕捉节点之间的关系和依赖性,在许多领域都有广泛应用,如社交网络分析、流行病学、化学分子预测等。下面我将介绍GNN的基本原理、数学公式和Python代码实现。

视频推荐：简单粗暴带你快速理解GNN

1. 基本原理

GNN的核心思想是通过迭代的方式聚合邻居节点的表示,从而更新当前节点的表示。每个节点的表示都是基于其自身特征和邻居节点的表示来更新的。这种传播机制使得GNN能够很好地捕捉图结构中节点之间的依赖关系。

GNN通常包括以下几个关键步骤:

1. 初始化节点表示
2. 邻居聚合(Neighbor Aggregation)
3. 节点更新(Node Update)
4. 输出表示(Output Representation)

2. 数学公式

假设有一个无向图 G=(V, E),其中 V 是节点集合,E 是边集合。对于每个节点 v∈V,我们用 $h_v^{(k)}$ 表示该节点在第 k 层的隐藏表示, $h_v^{(0)}$ 表示该节点的初始特征。

邻居聚合步骤可以表示为:

$$h_N(v{)}^{(k)}=AG{G}^{(k)}(h_u^{(k-1)},\forall u\in N(v))$$

其中 N(v) 表示节点 v 的邻居集合,AGG^{(k)} 表示在第 k 层使用的聚合函数。

节点更新步骤可以表示为:

$$h_v^{(k)}=UPDAT{E}^{(k)}(h_v^{(k-1)},h_N(v{)}^{(k)})$$

其中 $UPDAT{E}^{(k)}$ 表示在第 k 层使用的更新函数。

通过 K 次迭代,我们可以得到每个节点的最终表示 $h_v^{(K)}$。

3. Python 代码

下面是一个使用PyTorch Geometric实现的简单GNN模型示例,用于节点分类任务:

import torch
from torch.nn import Linear
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

# 定义GCN模型
class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

# 加载数据
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='./data/Cora', name='Cora')

# 获取数据
data = dataset[0]

# 模型实例化
model = GCN(dataset.num_node_features, 16, dataset.num_classes)

# 训练和评估模型
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = torch.nn.NLLLoss()

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

def test():
    model.eval()
    logits = model(data.x, data.edge_index)
    accs = []
    for _, mask in data('train_mask', 'val_mask', 'test_mask'):
        pred = logits[mask].max(1)[1]
        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()
        accs.append(acc)
    return accs

for epoch in range(200):
    loss = train()
    train_acc, val_acc, test_acc = test()
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')

# 保存训练好的模型参数
torch.save(model.state_dict(), 'your_model.pth')
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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

在这个示例中,我们使用了PyTorch Geometric库中的GCNConv层来实现图卷积操作。我们定义了一个两层GCN模型,用于在Cora数据集上进行节点分类任务。

运行这段代码,您将看到每个epoch的训练损失和准确率输出。请注意,这只是一个简单的GNN示例,在实际应用中,您可能需要根据具体问题和数据集调整模型结构和超参数。

简单使用模型

上述代码如果对不熟悉cora数据集的友友会有些困难，下面把数据结构直接展现出来，方便理解

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.data import Data

# 定义GCN模型
class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)



# 生成节点特征
node_features = torch.tensor([[1.0, 0.5, 0.2],
                               [0.4, 0.7, 0.1],
                               [0.3, 0.6, 0.8],
                               [0.6, 0.2, 0.4],
                               [0.5, 0.4, 0.7],
                               [0.2, 0.3, 0.6],
                               [0.8, 0.1, 0.3],
                               [0.1, 0.9, 0.2],
                               [0.7, 0.2, 0.5],
                               [0.3, 0.4, 0.6]], dtype=torch.float)

# 生成边索引
edge_index = torch.tensor([[0, 1, 2, 0, 3, 4, 5, 6, 6, 7, 8, 9, 1, 2, 3, 4, 5, 0, 1, 9],
                            [1, 2, 0, 3, 0, 5, 6, 6, 7, 8, 9, 1, 3, 4, 5, 0, 1, 9, 2, 8]], dtype=torch.long)

# 生成节点标签
node_labels = torch.tensor([0, 1, 0, 1, 2, 2, 0, 1, 2, 0])

# 创建PyG Data对象
data = Data(x=node_features, edge_index=edge_index, y=node_labels)

model = GCN(3, 16, 3)  # 输入特征维度为3,隐藏层维度为16,输出维度为3(对应3个类别)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.NLLLoss()

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out, data.y)
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    pred = model(data.x, data.edge_index).argmax(dim=1)
    correct = (pred == data.y).sum()
    acc = int(correct) / len(data.y)
    print(f'Accuracy: {acc:.4f}')

# 保存训练好的模型参数
torch.save(model.state_dict(), 'your_model.pth')



# 新的数据集
new_node_features = torch.tensor([[0.2, 0.3, 0.4],
                                  [0.5, 0.1, 0.7],
                                  [0.8, 0.2, 0.6]], dtype=torch.float)

new_edge_index = torch.tensor([[0, 1, 2],
                               [1, 2, 0]], dtype=torch.long)

# 创建新的PyG Data对象
new_data = Data(x=new_node_features, edge_index=new_edge_index)

# 加载训练好的模型
model = GCN(3, 16, 3)
model.load_state_dict(torch.load('your_model.pth'))
model.eval()

# 使用模型进行预测
with torch.no_grad():
    pred = model(new_data.x, new_data.edge_index).argmax(dim=1)

# 输出预测结果
print("Predictions:", pred)

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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

Accuracy: 0.7000
Predictions: tensor([1, 0, 0])
1
2