Pytorch Geometric - 1. 介绍

原文地址：INTRODUCTION BY EXAMPLE

博客地址：陈小默的CSDN

通过自带的几个样例，可以大致了解到Pytorch Gemetric的基本概念。

Graph数据的处理

Graph数据模型通常成对的包含两部分信息，即以边(edge)的形式表示的关系(relation) 和 以节点(node)的形式表示的对象。在Geometric中，描述graph对象的类是torch_geometric.data.Data，这个类包含以下属性:

• x: 描述节点特征的矩阵，形状为[num_nodes, num_node_features]
• edge_index: graph的边矩阵，形状[2, num_edges]，类型为torch.long
• edge_attr: 用于描述边特征的矩阵，形状为[num_edges, num_edge_features]
• y: 节点标签，对于node-level级的任务，该数据形状为[num_nodes, ]；对于graph-level级的任务，形状为[, ]
• pos: 该矩阵表示节点在空间中的位置信息。形状为[num_nodes, dimensions]

给定如下结构的graph:

在程序中的表示方式如下:

import torch
from torch_geometric.data import Data

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)
>>> Data(edge_index=[2, 4], x=[3, 1])

这里需要注意的是edge_index，第一行表示边的source节点，第二行对应位置为边的target节点。仅接受形状为[2, num_edges]的边信息，如果形状错误，需要将数据转置。

import torch
from torch_geometric.data import Data

edge_index = torch.tensor([[0, 1],
                           [1, 0],
                           [1, 2],
                           [2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index.t().contiguous())
>>> Data(edge_index=[2, 4], x=[3, 1])

注意：对于无向图，两个方向的边都需要写入。

对象创建完成后可以通过其中的属性查看对象信息：

print(data.keys)
>>> ['x', 'edge_index']

print(data['x'])
>>> tensor([[-1.0],
            [0.0],
            [1.0]])

for key, item in data:
    print("{} found in data".format(key))
>>> x found in data
>>> edge_index found in data

'edge_attr' in data
>>> False

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

# 把data对象移动到GPU.
device = torch.device('cuda')
data = data.to(device)

通用的基准数据集

Pytorch Geometric中包含大量的常见基准数据集。在初始化数据集的时候，框架会自动下载数据集的原始文件，并将其处理为Data对象。例如要下载ENZYMES数据集（由600个graph划分为6个类别）

from torch_geometric.datasets import TUDataset

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
>>> ENZYMES(600)

len(dataset)
>>> 600

dataset.num_classes
>>> 6

dataset.num_node_features
>>> 3

data = dataset[0]
>>> Data(edge_index=[2, 168], x=[37, 3], y=[1])

data.is_undirected()
>>> True

另一个数据集Cora，用于半监督图形分类的基准数据集：

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')
>>> Cora()

len(dataset)
>>> 1

dataset.num_classes
>>> 7

dataset.num_node_features
>>> 1433

data = dataset[0]
>>> 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.sum().item()
>>> 140

data.val_mask.sum().item()
>>> 500

data.test_mask.sum().item()
>>> 1000

Mini-batches

神经网络通常以batch的方式进行训练，geometric在mini-batch实现了并行化，这种组合允许在一个batch中使用不同数量的边和节点。

在torch_geometric.data.DataLoader中，已经包含了此过程。

from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in loader:
    batch
    >>> Batch(batch=[1082], edge_index=[2, 4066], x=[1082, 21], y=[32])

    batch.num_graphs
    >>> 32

torch_geometric.data.Batch继承自torch_geometric.data.Data，并包含一个名为batch的附加属性。batch是一个列向量，它将每个节点映射到该batch中的对应的graph：
$$\text{batch}=[0\dots 01\dots n-2n-1\dots n-1{]}^T$$
求平均节点特征的示例如下：

from torch_scatter import scatter_mean
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for data in loader:
    data
    >>> Batch(batch=[1082], edge_index=[2, 4066], x=[1082, 21], y=[32])

    data.num_graphs
    >>> 32

    x = scatter_mean(data.x, data.batch, dim=0)
    x.size()
    >>> torch.Size([32, 21])

Graph的学习方法

使用下列方法实现并训练一个两层的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 = 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