PyTorch Geometric入门指南
pytorch geometric是基于pytorch框架封装的图神经网络库,相比于DGL库更加便利,符合pytorch开发规范,拥有pytorch基础,更易于使用。pytorch基础请移步pytorch训练CNN
1.数据类型介绍
- 类比于pytorch,pyg拥有同样名称的库用于构建网络和处理数据:
库名 | 作用 |
---|---|
torch_geometric.nn | 封装了常用的GNN相关网络,直接可用 |
torch_geometric.data | GNN相关的数据结构,数据批量加载工具等 |
torch_geometric.datasets | 公开直接可用的数据集 |
torch_geometric.transforms | 数据转换工具类 |
torch_geometric.utils | 其他工具类 |
torch_geometric.io | 读取写入类 |
-
PyTorch Geometric中的一个图由的实例描述torch_geometric.data.Data,默认情况下具有以下属性:
- data.x:具有形状的节点特征矩阵 [num_nodes, num_node_features]
例如上图,有三个节点,每个节点有两个属性,则x大小为[3,2]
- data.edge_index: 具有形状和类型的COO格式的图形连接[2, num_edges]torch.long
图结构都有邻接关系,一般使用邻接矩阵,为了方便计算,pyg采用COO为默认形式:
即对邻接矩阵按顺序读取非零元素,记录其行列值;
- data.edge_attr: 具有形状的边缘特征矩阵 [num_edges, num_edge_features]
data.edge_index表示无权情况的邻接矩阵,但是当图是有权值或双向表示时,需要映射存储相应边的权值,即data.edge_attr;
- data.y: 要训练的目标(可以具有任意形状),例如,形状的节点级目标或形状的图形级目标[num_nodes, *][1, *]
表示训练之后结果矩阵,例如最常用的分为几类,one-hot编码形式
- data.pos: 具有形状的节点位置矩阵 [num_nodes, num_dimensions]
此属性针对具有位置信息的节点,例如点云数据
- data.x:具有形状的节点特征矩阵 [num_nodes, num_node_features]
2.训练过程
- 过程与pytorch训练其他网络,结构一模一样,例如官方教程给的例子,可对比CNN结构:
- 加载Cora数据集:
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
>>> Cora()
- 构建网络,现在实现一个两层的GCN等图网络:
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = GCNConv(dataset.num_node_features, 16)
self.conv2 = GCNConv(16, dataset.num_classes)
def forward(self, data):
x, edge_index = data.x, data.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)
- CPU上训练200次;(其中mask指掩膜映射,train代表训练集,test代表测试集,此示例数据集自带,实际使用可自定义或直接拆分使用)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
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()
- 测试
_, pred = model(data).max(dim=1)
correct = int(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / int(data.test_mask.sum())
print('Accuracy: {:.4f}'.format(acc))
>>> Accuracy: 0.8150