实现第一个图神经网络
1 基本概念
-
图数据的信息包含3个层面,分别是节点信息(V)、边信息(E)、图整体(U)信息,它们通常是用向量来表示。而图神经网络就是通过学习数据从而得到3个层面向量的最优表示。
-
对于图数据而言的任务分类:
- 图层面:分子是天然的图,原子是节点,化学键是边。现在要做一个分类,有一个苯环的分子分一类,两个苯环的分子分一类。
- 边层面:拳击赛上通过语义分割将台上的人和环境分离开,赛场上的人是节点,要做预测,预测这些人之间的关系,是对抗?还是观众?
- 节点层面:拳击赛上分成了两队,所有会员都是节点,哪些会员是A队,哪些是B队。
- 工作原理
GNN是对图上的所有属性进行的一个可以优化的变换,它的输入是一个图,输出也是个图。它只对属性向量(即上文所述的V、E、U)进行变换,但它不会改变图的连接性(即哪些点互相连接经过GNN后是不会变的)。在获取优化后的属性向量之后,再根据实际的任务,后接全连接神经网络,进行分类和回归。可以把图神经网络看做是一个图数据的在三个维度的特征提取器。
GNN对属性向量优化的方法叫做消息传递机制。比如最原始的GNN是SUM求和传递机制;到后面发展成图卷积网络(GCN)就考虑到了节点的度,度越大,权重越小,使用了加权的SUM;再到后面发展为图注意力网络GAT,在消息传递过程中引入了注意力机制;目前的SOTA模型研究也都专注在了消息传递机制的研究。见下图所示。
2 代码实现
常用的包是PyG(PyTorch Geometric),它是一个为图形数据的处理和学习提供支持的PyTorch扩展库,提供了一系列工具来帮助开发者轻松地实现基于图形的机器学习任务,例如图分类、图回归、图生成等。
PyG有许多内置的图分类和图回归数据集,可以用于训练和评估图神经网络。以下是一些常用的内置数据集:
- Cora, Citeseer, Pubmed:这些数据集是文献引用网络数据集,用于节点分类任务。
- PPI:蛋白质蛋白相互作用网络数据集,用于边分类任务。
- Reddit:Reddit社交网络数据集,用于节点分类任务。
- Amazon-Computers,Amazon-Photo:Amazon商品共同购买网络数据集,用于节点分类和图分类任务。
- ENZYMES:蛋白质分子结构数据集,用于图分类任务。
- MUTAG:分子化合物数据集,用于图分类任务。
- QM7b:有机分子数据集,用于图回归任务。
下面我使用PyG的内置数据进行3个任务的代码实现:
2.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%
2.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层数太少,可以进一步添加。
2.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 F.log_softmax(x, dim=-1)
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 test_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}')
3 总结
综合上面所有的内容,最重要的是以下两点:
①不同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)
图特征融合之后再跟几个线性层