图神经网络(Graph Neural Networks, GNNs)是一种专门处理图结构数据的深度学习模型。它可以很好地捕捉节点之间的关系和依赖性,在许多领域都有广泛应用,如社交网络分析、流行病学、化学分子预测等。下面我将介绍GNN的基本原理、数学公式和Python代码实现。
视频推荐:简单粗暴带你快速理解GNN
1. 基本原理
GNN的核心思想是通过迭代的方式聚合邻居节点的表示,从而更新当前节点的表示。每个节点的表示都是基于其自身特征和邻居节点的表示来更新的。这种传播机制使得GNN能够很好地捕捉图结构中节点之间的依赖关系。
GNN通常包括以下几个关键步骤:
- 初始化节点表示
- 邻居聚合(Neighbor Aggregation)
- 节点更新(Node Update)
- 输出表示(Output Representation)
2. 数学公式
假设有一个无向图 G=(V, E),其中 V 是节点集合,E 是边集合。对于每个节点 v∈V,我们用 h v ( k ) h_v^{(k)} hv(k) 表示该节点在第 k 层的隐藏表示, h v ( 0 ) h_v^{(0)} hv(0) 表示该节点的初始特征。
邻居聚合步骤可以表示为:
h N ( v ) ( k ) = A G G ( k ) ( h u ( k − 1 ) , ∀ u ∈ N ( v ) ) h_N(v)^{(k)} = AGG^{(k)}({h_u^{(k-1)}, ∀u∈N(v)}) hN(v)(k)=AGG(k)(hu(k−1),∀u∈N(v))
其中 N(v) 表示节点 v 的邻居集合,AGG^{(k)} 表示在第 k 层使用的聚合函数。
节点更新步骤可以表示为:
h v ( k ) = U P D A T E ( k ) ( h v ( k − 1 ) , h N ( v ) ( k ) ) h_v^{(k)} = UPDATE^{(k)}(h_v^{(k-1)}, h_N(v)^{(k)}) hv(k)=UPDATE(k)(hv(k−1),hN(v)(k))
其中 U P D A T E ( k ) UPDATE^{(k)} UPDATE(k) 表示在第 k 层使用的更新函数。
通过 K 次迭代,我们可以得到每个节点的最终表示 h v ( K ) h_v^{(K)} hv(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')
在这个示例中,我们使用了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)
Accuracy: 0.7000
Predictions: tensor([1, 0, 0])