【PyG 教程】PyTorch Geometric 安装与入门
文章作者：梦家
个人站点：dreamhomes.top
原文地址：https://dreamhomes.github.io/posts/202005182019.html
公众号ID：DreamHub
早期基于 DGL 库学习卷积神经网络，写过一个 GCN demo。后来PyTorch的几何扩展库出来了，发现学术界很多paper都是基于 PyG 实现的，因此学习下 PyG 如何使用。
事实上这两个库都非常实用，但 PyG 和 DGL 这两大框架应该如何选择呢？没有好坏之分，个人只是从工具生态中进行判断，给出这两个库在Github中 Fork 和 Star 数量，可以说明 PyG 维护人员和受欢迎的程度还是要高一点的啊！
PyG 安装
PyG 全称是PyTorch-Geometric，是一个PyTorch基础上的一个库，专门用于图形式的数据，可以加速图学习算法的计算过程，比如稀疏化的图等。
参考：https://github.com/rusty1s/pytorch_geometric
对于 PyTorch 版本为1.5.0，安装一下库：
$ pip install torch-scatter==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.5.0.html
$ pip install torch-sparse==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.5.0.html
$ pip install torch-cluster==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.5.0.html
$ pip install torch-spline-conv==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.5.0.html
$ pip install torch-geometric
其中${CUDA}取决于为本地的PyTorch安装环境，可以替换为：cpu, cu92, cu101 or cu102。
如果以上命令执行完成说明 torch_geometric 已安装成功！
问题
理论上安装完成后没问题但是我导入 torch_geometric 时报错：
RuntimeError:
Tried to access nonexistent attribute or method 'true_divide_' of type 'Tensor'.:
File "/usr/local/lib/python3.6/dist-packages/torch_scatter/scatter.py", line 53
count = broadcast(count, out, dim)
if torch.is_floating_point(out):
out.true_divide_(count)
~~~~~~~~~~~~~~~~ <--- HERE
else:
out.floor_divide_(count)
问题待解决，已提交issue：https://github.com/rusty1s/pytorch_scatter/issues/140
入门实例
参考官方教程：https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html
本文采用官方的实例进行入门，主要包括：
- Data Handling of Graphs：图数据处理
- Common Benchmark Datasets：公开的基准数据集
- Mini-batches：小批量数据
- Data Transforms：数据变换
- Learning Methods on Graphs：图上的学习方法
以下实验部分在jupyter中学习更加哦！传送门~
Data Handling of Graphs
图是用来建模现实系统中实体（节点）和关系（边）的，简单图在 PyG 中通常存储为实例：torch_geometric.data.Data 实例，具有一下属性值：
data.x: 节点特征矩阵，大小为[num_nodes, num_node_features]data.edge_index: COO格式的边存储形式，大小为[2, num_edges]data.edge_attr: 边属性，大小为[num_edges, num_edge_features]data.y: 任意大小的训练分类目标, 例如，节点级别大小为[num_nodes, *]或图级别大小为[1, *]data.pos: 节点的位置矩阵，大小为[num_nodes, num_dimensions]
以代码示例说明上述属性，对于三个节点四条边的无向无权图：
import torch
from torch_geometric.data import Data
# node index can start without 0
edge_index = torch.tensor([[0, 1, 1, 2],
[1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
data = Data(x=x, edge_index=edge_index)
注意：如果边不是以COO形式给出的而是以节点对的形式，需要先转置t()再利用函数contiguous()，例如以上示例：
edge_index = torch.tensor([[0,1], [1,0], [1,2], [2,1]], dtype=torch.long)
edge_index.t().contiguous()
tensor([[0, 1, 1, 2],
[1, 0, 2, 1]])
data # 2, 4
Data(edge_index=[2, 4], x=[3, 1])
data.x
tensor([[-1.],
[ 0.],
[ 1.]])
data.y
data.edge_index
tensor([[0, 1, 1, 2],
[1, 0, 2, 1]])
data.edge_attr
data.pos
for key, item in data:
print(key, item)
edge_index tensor([[0, 1, 1, 2],
[1, 0, 2, 1]])
x tensor([[-1.],
[ 0.],
[ 1.]])
data.num_nodes
3
data.num_edges
4
data.num_node_features
1
data.contains_isolated_nodes()
False
data.contains_self_loops()
False
data.is_directed()
False
# transfer data to GPU
if torch.cuda.is_available():
data = data.to(torch.device('cuda'))
Common Benchmark Datasets
- Planetoid datasets
- Cora
- Citeseer
- Pubmed
- Graph classification datasets.
- QM7 and QM9 datasets.
- 3D mesh/point
TUDataset
from torch_geometric.datasets import TUDataset
dataset = TUDataset(root='./data', name='ENZYMES')
len(dataset) # 600 graphs for graph classification.
600
dataset.num_classes # 6 classes
6
dataset.num_node_features
3
dataset.num_edge_attributes
0
data = dataset[0] # 37 nodes with 3 features,168/2 edges and the graph is assigned to clase 1.
data
Data(edge_index=[2, 168], x=[37, 3], y=[1])
data.is_directed()
False
Split the dataset:90/10 train/test
train_dataset = dataset[:540]
test_dataset = dataset[540:]
train_dataset
ENZYMES(540)
Dataset permutation
dataset = dataset.shuffle() # dataset = dataset[torch.randperm(len(dataset))]
dataset[0]
Data(edge_index=[2, 104], x=[30, 3], y=[1])
Cora dataset
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='./data', name='Cora')
len(dataset)
1
data = dataset[0]
data
Data(edge_index=[2, 10556], test_mask=[2708], train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])
data.is_undirected()
True
data.train_mask
tensor([ True, True, True, ..., False, False, False])
data.train_mask.sum().item()
140
data.val_mask.sum().item()
500
data.test_mask.sum().item()
1000
data.y
tensor([3, 4, 4, ..., 3, 3, 3])
train_mask训练的节点val_mask验证节点test_mask测试节点
Mini-batchs
神经网络通常会按照batch方式进行训练，PyG 通过构建稀疏化的分块对角阵实现mini-batch的并行化，构建方式按照每一个Data实例的edge_index构建一个图的邻接矩阵，然后将所有节点的特征向量按行拼接。使得不同数量的顶点数和边数的图可以一起训练。可以使用PyG内部的 torch_geometric.data.DataLoader进行图拼接的过程！
A = [ A 1 ⋱ A n ] , X = [ X 1 ⋮ X n ] , Y = [ Y 1 ⋮ Y n ] \mathbf{A}=\left[\begin{array}{ccc} \mathbf{A}_{1} & & \\ & \ddots & \\ & & \mathbf{A}_{n} \end{array}\right], \quad \mathbf{X}=\left[\begin{array}{c} \mathbf{X}_{1} \\ \vdots \\ \mathbf{X}_{n} \end{array}\right], \quad \mathbf{Y}=\left[\begin{array}{c} \mathbf{Y}_{1} \\ \vdots \\ \mathbf{Y}_{n} \end{array}\right] A=⎣⎡A1⋱An⎦⎤,X=⎣⎢⎡X1⋮Xn⎦⎥⎤,Y=⎣⎢⎡Y1⋮Yn⎦⎥⎤
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader
dataset = TUDataset(root='./data', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
len(loader)
19
for batch in loader:
print(batch)
Batch(batch=[1034], edge_index=[2, 3968], x=[1034, 21], y=[32])
Batch(batch=[1057], edge_index=[2, 4196], x=[1057, 21], y=[32])
Batch(batch=[1034], edge_index=[2, 4078], x=[1034, 21], y=[32])
Batch(batch=[1167], edge_index=[2, 4186], x=[1167, 21], y=[32])
Batch(batch=[1096], edge_index=[2, 4132], x=[1096, 21], y=[32])
Batch(batch=[1176], edge_index=[2, 4286], x=[1176, 21], y=[32])
Batch(batch=[1070], edge_index=[2, 4226], x=[1070, 21], y=[32])
Batch(batch=[1077], edge_index=[2, 4114], x=[1077, 21], y=[32])
Batch(batch=[1085], edge_index=[2, 4276], x=[1085, 21], y=[32])
Batch(batch=[949], edge_index=[2, 3546], x=[949, 21], y=[32])
Batch(batch=[1049], edge_index=[2, 3944], x=[1049, 21], y=[32])
Batch(batch=[940], edge_index=[2, 3690], x=[940, 21], y=[32])
Batch(batch=[1298], edge_index=[2, 4444], x=[1298, 21], y=[32])
Batch(batch=[935], edge_index=[2, 3650], x=[935, 21], y=[32])
Batch(batch=[982], edge_index=[2, 3866], x=[982, 21], y=[32])
Batch(batch=[1022], edge_index=[2, 3908], x=[1022, 21], y=[32])
Batch(batch=[916], edge_index=[2, 3514], x=[916, 21], y=[32])
Batch(batch=[982], edge_index=[2, 3840], x=[982, 21], y=[32])
Batch(batch=[711], edge_index=[2, 2700], x=[711, 21], y=[24])
batch大小为32个图，但是每一个图的规模是不一样的，上例中第一个batch内的32个图共1005节点，含有3948条边。
torch_geometric.data.Batch继承torch_geometric.data.Data，并且添加了一个额外的属性batch。batch是一个列向量，代表了每一个节点对应到哪一个图。
batch = [ 0 ⋯ 0 1 ⋯ n − 2 n − 1 ⋯ n − 1 ] ⊤ \text { batch }=\left[\begin{array}{cccccccc} 0 & \cdots & 0 & 1 & \cdots & n-2 & n-1 & \cdots & n-1 \end{array}\right]^{\top} batch =[0⋯01⋯n−2n−1⋯n−1]⊤
可以根据Batch对每个图中的节点特征进行平均化，其中使用到 scatter 库，以每一个图为单位，将各个图中的所有节点的特征向量计算了一个平均值，所以维度为[32, 21]。
from torch_scatter import scatter_mean
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader
dataset = TUDataset(root='./data', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for data in loader:
print(data)
print(data.num_graphs)
x = scatter_mean(data.x, data.batch, dim=0)
print(x.size())
Batch(batch=[971], edge_index=[2, 3762], x=[971, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1214], edge_index=[2, 4626], x=[1214, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1102], edge_index=[2, 4184], x=[1102, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1076], edge_index=[2, 3630], x=[1076, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1044], edge_index=[2, 3978], x=[1044, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1152], edge_index=[2, 4216], x=[1152, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1053], edge_index=[2, 4152], x=[1053, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1018], edge_index=[2, 3876], x=[1018, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1076], edge_index=[2, 4082], x=[1076, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[941], edge_index=[2, 3586], x=[941, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1119], edge_index=[2, 4292], x=[1119, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1026], edge_index=[2, 3818], x=[1026, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[934], edge_index=[2, 3774], x=[934, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1036], edge_index=[2, 3882], x=[1036, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[983], edge_index=[2, 3734], x=[983, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1003], edge_index=[2, 3918], x=[1003, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[1140], edge_index=[2, 4506], x=[1140, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[924], edge_index=[2, 3572], x=[924, 21], y=[32])
32
torch.Size([32, 21])
Batch(batch=[768], edge_index=[2, 2976], x=[768, 21], y=[24])
24
torch.Size([24, 21])
Data Transforms
PyG 中的数据变换与 torchvision 中图片变换与扩充类似。变换操作可以是使用torch_geometric.transform.Compose进行图拼接！
下面以ShapeNet（17000 3D点云，16种形状） 数据集为例，下载数据时可以根据 kNN 对点云进行图构造：
from torch_geometric.datasets import ShapeNet
dataset = ShapeNet(root='./data/ShapeNet', categories=['Airplane'])
dataset
ShapeNet(2349, categories=['Airplane'])
dataset[0]
Data(category=[1], pos=[2518, 3], x=[2518, 3], y=[2518])
可以根据KNN构造点云图，下载时已预处理下次使用时不要在进行图构造！
import torch_geometric.transforms as T
dataset = ShapeNet(root='./data/ShapeNet', categories=['Airplane'],
pre_transform=T.KNNGraph(k=6),
transform=T.RandomTranslate(0.01))
dataset[0]
D:\DevTools\Miniconda3\lib\site-packages\torch_geometric\data\dataset.py:151: UserWarning: The `pre_transform` argument differs from the one used in the pre-processed version of this dataset. If you really want to make use of another pre-processing technique, make sure to delete `data\ShapeNet\processed` first.
'sure to delete `{}` first.'.format(self.processed_dir))
Data(category=[1], pos=[2518, 3], x=[2518, 3], y=[2518])
Learning Methods on Graphs
在学习PyG的数据处理与图变换方法后，接下来就依赖 PyG 实现 GNN了！
# 加载数据集
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='./data', name='Cora')
dataset
Cora()
# 2-layers 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)
# 训练
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)
model.train()
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()
# 测试
model.eval()
_, pred = model(data).max(dim=1)
correct = float (pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))
Accuracy: 0.7970
本文简单说明的PyG的安装过程并以简单的实例介绍了 PyG 中的基础知识。接下来需要学习如何构造自定义的图卷积操作！