图神经网络（GNN）最简单全面原理与代码实现

最新推荐文章于 2025-01-31 13:00:05 发布

hellozhxy

最新推荐文章于 2025-01-31 13:00:05 发布

阅读量2.4w

点赞数 85

文章标签： GNN

原文链接：https://zhuanlan.zhihu.com/p/603486955

版权

GNN，即图神经网络，是一种用于处理图形数据的深度学习技术。

1. 什么是图数据？在图神经网络中，图数据是以什么形式表示的？

图数据是由节点（Node）和边（Edge）组成的数据，最简单的方式是使用邻接矩阵来表示图形结构，从而捕捉图形中的节点和边的相关性。假设图中的节点数为n，那么邻接矩阵就是一个n*n的矩阵，如果节点之间有关联，则在邻接矩阵中表示为1，无关联则为0。在图中，鲁班与其他英雄都没有关联，表现在邻接矩阵当中就是它所在的行与列为全零。

王者荣耀当中的图和邻接矩阵

图数据的信息包含3个层面，分别是节点信息（V）、边信息（E）、图整体（U）信息，它们通常是用向量来表示。而图神经网络就是通过学习数据从而得到3个层面向量的最优表示。

2. 对于图数据而言有怎样的任务？

●图层面的任务（分类/回归）

例：分子是天然的图，原子是节点，化学键是边。现在要做一个分类，有一个苯环的分子分一类，两个苯环的分子分一类。这是图分类任务。

●边层面的任务（分类/回归）

例：UFO拳击赛上，首先通过语义分割把台上的人和环境分离开来。赛场上的人都是节点，现在要做一个预测，预测的是这些人之间的关系，是对抗关系？还是观众watch的关系？还是裁判watch的关系？这是边分类任务。

●节点层面的任务（分类/回归）

例：假设一个跆拳道俱乐部里有A、B两个教练，所有的会员都是节点。有一天A、B两个跆拳道教练决裂，那么各个学员是愿意和A在一个阵营还是愿意和B在一个阵营？这是节点分类任务。

3. 图神经网络是如何工作的？

GNN工作流程图

GNN是对图上的所有属性进行的一个可以优化的变换，它的输入是一个图，输出也是个图。它只对属性向量（即上文所述的V、E、U）进行变换，但它不会改变图的连接性（即哪些点互相连接经过GNN后是不会变的）。在获取优化后的属性向量之后，再根据实际的任务，后接全连接神经网络，进行分类和回归。大家可以把图神经网络看做是一个图数据的在三个维度的特征提取器。

GNN对属性向量优化的方法叫做消息传递机制。比如最原始的GNN是SUM求和传递机制；到后面发展成图卷积网络（GCN）就考虑到了节点的度，度越大，权重越小，使用了加权的SUM；再到后面发展为图注意力网络GAT，在消息传递过程中引入了注意力机制；目前的SOTA模型研究也都专注在了消息传递机制的研究。见下图所示。

三种不同的图神经网络模型的消息传递机制差异

但是！即使消息传递机制你不完全明白也没有关系，你只要记住：不同GNN的本质差别就在于它们如何进行节点之间的信息传递和计算，也就是它们的消息传递机制不同。就可以了！

4. 图神经网络代码实现

我常用的包是PyG（PyTorch Geometric），它是一个为图形数据的处理和学习提供支持的PyTorch扩展库，提供了一系列工具来帮助开发者轻松地实现基于图形的机器学习任务，例如图分类、图回归、图生成等。

PyG有许多内置的图分类和图回归数据集，可以用于训练和评估图神经网络。以下是一些常用的内置数据集：

Cora, Citeseer, Pubmed：这些数据集是文献引用网络数据集，用于节点分类任务。
PPI：蛋白质蛋白相互作用网络数据集，用于边分类任务。
Reddit：Reddit社交网络数据集，用于节点分类任务。
Amazon-Computers，Amazon-Photo：Amazon商品共同购买网络数据集，用于节点分类和图分类任务。
ENZYMES：蛋白质分子结构数据集，用于图分类任务。
MUTAG：分子化合物数据集，用于图分类任务。
QM7b：有机分子数据集，用于图回归任务。

下面我将使用PyG的内置数据进行3个任务的代码实现：

4.1 节点分类任务代码实现

Cora数据集是PyG内置的节点分类数据集，代表着学术论文的相关性分类问题（即把每一篇学术论文都看成是节点），Cora数据集有2708个节点，1433维特征，边数为5429。标签是文献的主题，共计 7 个类别。所以这是一个7分类问题。

下面是代码【需要使用美国的IP，否则好像不能下载Cora数据，大伙可以试试】：

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
#载入数据
dataset = Planetoid(root='~/tmp/Cora', name='Cora')
data = dataset[0]
#定义网络架构
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(dataset.num_features, 16)  #输入=节点特征维度，16是中间隐藏神经元个数
        self.conv2 = GCNConv(16, dataset.num_classes)
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
#模型训练
model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)    #模型的输入有节点特征还有边特征,使用的是全部数据
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])   #损失仅仅计算的是训练集的损失
    loss.backward()
    optimizer.step()
#测试：
model.eval()
test_predict = model(data.x, data.edge_index)[data.test_mask]
max_index = torch.argmax(test_predict, dim=1)
test_true = data.y[data.test_mask]
correct = 0
for i in range(len(max_index)):
    if max_index[i] == test_true[i]:
        correct += 1
print('测试集准确率为：{}%'.format(correct*100/len(test_true)))

对于这个节点7分类的问题，最终在测试集（1000个样本）上的分类准确率为79.9%（见下图）。因为我们只是使用了一个很简单的模型架构，所以这个结果还说得过去。

测试结果

4.2 边分类任务代码实现

同样是利用Cora数据集，只是这个时候我们关注的不再是节点特征，而是边特征，因此，在这里我们需要手动创建边标签的正例与负例。这是一个二分类问题。

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
from torch_geometric.utils import negative_sampling

# 边分类模型
class EdgeClassifier(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(EdgeClassifier, self).__init__()
        self.conv = GCNConv(in_channels, out_channels)
        self.classifier = torch.nn.Linear(2 * out_channels, 2)  

    def forward(self, x, edge_index):
        x = F.relu(self.conv(x, edge_index))
        pos_edge_index = edge_index    
        total_edge_index = torch.cat([pos_edge_index, negative_sampling(edge_index, num_neg_samples=pos_edge_index.size(1))], dim=1
        edge_features = torch.cat([x[total_edge_index[0]], x[total_edge_index[1]]], dim=1)  
        return self.classifier(edge_features)

# 加载数据集
dataset = Planetoid(root='./data/Cora/raw', name='Cora')
data = dataset[0]

# 创建train_mask和test_mask
edges = data.edge_index.t().cpu().numpy()   
num_edges = edges.shape[0]
train_mask = torch.zeros(num_edges, dtype=torch.bool)
test_mask = torch.zeros(num_edges, dtype=torch.bool)
train_size = int(0.8 * num_edges)
train_indices = torch.randperm(num_edges)[:train_size]
train_mask[train_indices] = True
test_mask[~train_mask] = True

# 定义模型和优化器/训练/测试
model = EdgeClassifier(dataset.num_features, 64)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

def train():
    model.train()
    optimizer.zero_grad()
    logits = model(data.x, data.edge_index)
    pos_edge_index = data.edge_index
    pos_labels = torch.ones(pos_edge_index.size(1), dtype=torch.long)  
    neg_labels = torch.zeros(pos_edge_index.size(1), dtype=torch.long)  
    labels = torch.cat([pos_labels, neg_labels], dim=0).to(logits.device)
    new_train_mask = torch.cat([train_mask, train_mask], dim=0)
    loss = F.cross_entropy(logits[new_train_mask], labels[new_train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

def test():
    model.eval()
    with torch.no_grad():
        logits = model(data.x, data.edge_index)
        pos_edge_index = data.edge_index
        pos_labels = torch.ones(pos_edge_index.size(1), dtype=torch.long)
        neg_labels = torch.zeros(pos_edge_index.size(1), dtype=torch.long)
        labels = torch.cat([pos_labels, neg_labels], dim=0).to(logits.device)
        new_test_mask = torch.cat([test_mask, test_mask], dim=0)
        
        predictions = logits[new_test_mask].max(1)[1]
        correct = predictions.eq(labels[new_test_mask]).sum().item()
        return correct / len(predictions)

for epoch in range(1, 1001):
    loss = train()
    acc = test()
    print(f"Epoch: {epoch:03d}, Loss: {loss:.4f}, Acc: {acc:.4f}")

数据流理解

在这里的mask部分，也许有的同学还没有完全理解。在这里着重解释：在创建模型时是根据所有的边创建正负样本。但是在训练过程当中，只取出train_mask的正负样本计算损失，对应于new_train_mask（new_train_mask = torch.cat([train_mask, train_mask], dim=0)），对于test亦然。

最终结果

最终在测试集上二分类准确率达到0.71。这个结果一般，这是因为模型架构过于简单。

①在计算边特征时，简单进行源节点特征和目标节点特征的concat。可以考虑其他方法（点乘等等），也可以在这里加MLP用以学习更多的节点-边模式

②GCN层数太少，可以进一步添加。

当然，在这里只是为了向大家展示GNN的简洁性，让大家能够最快地理解边分类的数据流，因此不做进一步的拓展。

4.3 图分类任务代码实现

在这里采用ENZYMES数据集。ENZYMES是一个常用的图分类基准数据集。它是由600个图组成的，这些图实际上表示了不同的蛋白酶的结构，这些蛋白酶分为6个类别（每个类别有100个蛋白酶）。因此，每个图代表一个蛋白酶，我们的任务是预测蛋白酶属于哪一个类别。这是6分类任务。

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader

# 加载数据集
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
dataset = dataset.shuffle()

train_dataset = dataset[:540]
test_dataset = dataset[540:]

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 定义图卷积网络模型
class GCN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(dataset.num_node_features, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels)
        self.lin = torch.nn.Linear(hidden_channels, dataset.num_classes)
    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = self.conv2(x, edge_index)
        x = x.relu()
        x = self.conv3(x, edge_index)
        x = global_mean_pool(x, batch)    # 使用全局平均池化获得图的嵌入
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin(x)
        return x

model = GCN(hidden_channels=64)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()

def train():
    model.train()
    for data in train_loader:
        optimizer.zero_grad()
        out = model(data.x, data.edge_index, data.batch)
        loss = criterion(out, data.y)
        loss.backward()
        optimizer.step()

def test(loader):
    model.eval()
    correct = 0
    for data in loader:
        out = model(data.x, data.edge_index, data.batch)
        pred = out.argmax(dim=1)
        correct += int((pred == data.y).sum())
    return correct / len(loader.dataset)

for epoch in range(1, 1001):
    train()
    train_acc = test(train_loader)
    test_acc = test(test_loader)
    print(f'Epoch: {epoch:03d}, Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')

这样就可以实现一个最简单的图分类任务了。

5 总结

综合上面所有的内容，最重要的是以下两点：

①不同GNN的本质区别是他们的消息传递机制不同，如GCN/GraphSAGE/GIN/GAT等等，只需要修改层的名称即可，目前已经达到了高度的集成化，不需要进行手撸，除非你的研究需要。

②三种不同的任务，他们的本质区别就是：Output层的输入不一样。

●对于节点层面的任务而言

可以直接self.conv = GCNConv(16, dataset.num_classes) ————这是直接把任务融合到卷积层

也可以在卷积获取特征之后，后面加几个线性层

●对于边层面的任务而言

通过GNN提取出节点信息，输入Output层之前需要进行边特征的融合（在这里是Concat节点特征）

边特征融合之后再跟几个线性层

edge_features = torch.cat([x[total_edge_index[0]], x[total_edge_index[1]]], dim=1)

●对于图层面的任务而言

通过GNN提取出节点信息，输入Output层之前需要进行图特征的融合（在这里是对节点特征进行全局平均池化）

 x = global_mean_pool(x, batch)

图特征融合之后再跟几个线性层

最后

如果你认真看完上述所有内容，你已经初步掌握了GNN的概念和使用方法。若想进阶，请使用自己的Graph_data进行尝试。