PyG

介绍

PyG（PyTorch Geometric）是一个基于PyTorch开发的图神经网络（GNN）专用库，旨在简化图结构数据的深度学习任务，能够与PyTorch无缝兼容。

PyG提供很多种GNN层，包括：GCNConv（图卷积）、GATConv（图注意力）、SAGEConv（GraphSAGE）、TransformerConv（图Transformer）等。

PyG内置Cora、Pubmed等经典数据集，支持一键下载与预处理。

安装与使用

1
pip install torch_geometric

图的数据表示

创建

PyG要求数据以Data对象传递，Data是PyG中表示单张图数据的容器类，继承自torch.Tensor，专门用于存储图结构数据。

Data的核心属性：

  1. 必需属性：

    属性名 数据类型 维度要求 说明
    x torch.Tensor [num_nodes, num_node_features] 节点特征矩阵（若没有节点特征，可不设置）
    edge_index torch.LongTensor [2, num_edges] 边的连接关系（COO 格式，第0行为源节点，第1行为目标节点）

  2. 可选属性：

    属性名 数据类型 说明
    edge_attr torch.Tensor 边特征矩阵，维度为 [num_edges, num_edge_features]
    y torch.Tensor 标签（节点级标签：[num_nodes]；图级标签：[1]）
    pos torch.Tensor 节点坐标（如 3D 点云数据），维度为 [num_nodes, 3]
    batch torch.LongTensor 批处理索引（用于区分多张图合并后的节点归属）
    自定义属性名 可以添加任意自定义的属性

比如：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch
from torch_geometric.data import Data

# 节点特征，4个节点，每个节点三个特征
x = torch.tensor([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
    [7.0, 8.0, 9.0],
    [10.0, 11.0, 12.0]
], dtype=torch.float)

# 边连接关系：3条边（0→1, 1→2, 2→3）
edge_index = torch.tensor([
    [0, 1, 2],
    [1, 2, 3]
], dtype=torch.long)

# 标签：节点分类任务（4个节点，2个类别）
y = torch.tensor([0, 1, 0, 1], dtype=torch.long)

# 创建Data对象
data = Data(x=x, edge_index=edge_index, y=y)

验证

数据验证：使用.validate()方法检查数据格式是否正确。

1
data.validate()	# 若数据合法，不抛出异常

实用方法

  1. 设备迁移

    1
    2
    3
    4
    # 将数据迁移到GPU
    data = data.to('cuda:0')
    # 检查设备
    print(data.x.device)

  2. 属性操作

    1
    2
    3
    4
    # 添加自定义属性
    data.custom_attr = "This is a test graph"
    # 查看所有属性
    print(data)	# 输出: Data(x=[4,3], edge_index=[2,3], y=[4], custom_attr='...')

  3. 图结构分析

    方法 说明
    data.num_nodes 返回节点数（若未显式设置 x，需通过 edge_index 推断）
    data.num_edges 返回边数（等于 edge_index.shape[1]）
    data.is_directed() 判断是否为有向图（若存在反向边则为无向图）

异构图的创建

异构图需要使用HeteroData：

1
2
import torch
from torch_geometric.data import HeteroData

定义节点类型与特征：

1
2
3
4
5
6
7
8
9
10
11
# 定义节点特征
user_features = torch.randn(10, 16)	# 10个用户节点，每个节点16维特征
product_features = torch.randn(20, 32)	# 20个商品节点，每个节点32维特征

# 创建HeteroData对象
data = HeteroData()

# 添加节点特征到HeteroData
data['user'].x = user_features
data['product'].x = product_features

添加边索引：使用三元组(源节点类型，边类型，目标节点类型)定义边，并分配COO格式的边索引（形状为[2, num_edges]）

1
2
3
4
5
6
7
8
9
10
11
# 用户到商品的“购买”关系边
edge_index = torch.tensor([
    [0, 1, 2],	# 源节点索引（用户）
    [0, 1, 0]	# 目标节点索引（商品）
])

data['user', 'buys', 'product'].edge_index = edge_index

# 用户与到商品的“浏览”关系边
edge_index_views = torch.tensor([[0, 2], [1, 0]])
data['user', 'views', 'product'] = edge_index_views

添加边特征（可选）：可以为边添加特征

1
data['user', 'buys', 'product'].edge_attr = torch.randn(3, 5)	# 每条边5个特征

验证异构图结构：打印HeteroData对象查看结构信息

1
print(data)

输出示例：

1
2
3
4
5
6
7
8
9
10
11
HeteroData(
  user={ x=[10, 16] },
  product={ x=[20, 32] },
  (user, buys, product)={
    edge_index=[2, 3],
    edge_attr=[3, 5]
  },
  (user, views, product)={
    edge_index=[2, 2]
  }
)

检查节点和边的数量：

1
2
3
4
5
6
7
# 节点数量
print("用户节点数:", data['user'].num_nodes)
print("商品节点数:", data['product'].num_nodes)

# 边数量
print("购买数量:", data['user', 'buys', 'product'].num_edges)
print("浏览边数:", data['user', 'views', 'product'].num_edges)

转换为GPU张量（可选）：

1
data = data.to('cuda:0')

注意：如果某类节点无特征，可以不设置.x属性，但某些模型可能需要初始化（例如使用全零张量）

模型构建

GCN

PyG中的GCNConv是GCN的核心层，

1
2
3
4
5
6
7
8
9
10
11
12
from torch_geometric.nn import GCNConv

# GCNConv参数说明
GCNConv(
    in_channels: int,				# 输入特征维度（节点特征数）
    out_channels: int,				# 输出特征维度
    improved: bool = False,			# 是否使用改进的邻接矩阵（邻接矩阵+2I，而不是I）
    cached: bool = False,			# 是否缓存归一化后的邻接矩阵（加速训练）
    add_self_loops: bool = True,	# 是否添加自环（默认添加）
    normalize: bool = True,			# 是否对邻接矩阵对称归一化（D^(-1/2)AD^(-1/2)）
    bias: bool = True,				# 是否添加偏置项
)

CGNConv()内部不包含激活函数，需手动添加。

定义一个两层的GCN：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch
import torch.nn as nn
from torch_geometric.nn import GCNConv

class GCN(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, dropout=0.3):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)
        self.dropout = nn.Dropout(dropout)
        self.relu = nn.ReLU()
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        
        x = self.conv1(x, edge_index)
        x = self.relu(x)
        x = self.dropout(x)
        
        x = self.conv2(x, edge_index)
        
        return x	# [num_nodes, out_channels]

RGCN

PyG中的RGCNConv是RGCN的核心层：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from torch_geometric.nn import RGCNConv

# RGCNConv参数说明
RGCNConv(
    in_channels: int,				# 输入特征维度（节点特征数）
    out_channels: int, 				# 输出特征维度
    num_relations: int, 			# 边类型总数（关系总数）
    num_bases: Optional[int]=None, 	# 基分解数量（用于减少多关系参数，与num_blocks互斥） 
    num_blocks: Optional[int]=None,	# 块分解数量（块对角矩阵优化，与num_bases互斥）
    aggr: str='mean',				# 聚合方式（'mean'/'sum'/'max'，默认'mean）
    bias: bool=True,				# 是否添加偏置项
    add_self_loops: bool=True		# 是否添加自环
    normalize: bool=True,			# 是否对邻接矩阵对称归一化（D^(-1/2)AD^(-1/2)）
    cached: bool=False				# 是否缓存归一化后的邻接矩阵
)

  • num_bases：基分解，通过共享基矩阵\(\mathbf{V}_b\)组合生成各种关系的\(\mathbf{W}_r\)： \[ \mathbf{W}_r = \sum_{b=1}^B a_{rb}\mathbf{V}_b \] 适用场景：关系间存在潜在共性（如社交网络中不同互动类型有相似模式）

  • num_blocks：块对角分解，将权重矩阵拆分为块对角结构，减少稠密参数 \[ \mathbf{W}_r = \mathrm{diag}(\mathbf{Q}_{1r},\cdots,\mathbf{Q}_{Br}) \] 适用场景：关系类型差异较大，需保持参数独立性

RGCNConv()对象输入数据格式：

  • 节点特征矩阵x，维度为[num_nodes, in_channels]（如果是异构图，有不同类型的节点，需要将它们合并）
  • 边索引edge_index，维度为[2, num_edges]（如果是异构图，有不同类型的边，需要将它们合并）
  • 边类型edge_type，维度为[num_edges]（值范围0到num_relations-1，第\(i\)个值代表上面边索引edge_index中第\(i\)条边的类型）

由此可以看出，RGCNConv()仅区分边类型，不区分节点类型。

RGCNConv()内部并不包含激活函数，需手动添加。

示例：

假设一个社交网络，包含3个用户，两种关秀（关注（类型0）、好友（类型1））

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import torch
import torch.nn as nn
from torch_geometric.nn import RGCNConv
from torch_geometric.data import Data

# 节点特征(3个节点，每个特征3维)
x = torch.randn(3, 3)

# 边索引（两种关系）
edge_index_0 = torch.tensor([
    [0, 2],
    [1, 1]
], dtype=torch.long)	# 关注关系，为有向图（0→1与2→1）
edge_index_1 = torch.tensor([
    [2, 1],
    [1, 2]
], dtype=torch.long)	# 好友关系，为无向图（2→1与1→2，表示为双向图）

# 合并边索引并标记类型
edge_index = torch.cat([edge_index_0, edge_index_1], dim=1)
edge_type = torch.tensor([0, 0, 1, 1], dtype=torch.long)	# 前两条边为关系0，后两条边为关系1

data = Data(
    x=x,
    edge_index=edge_index,
    edge_type=edge_type
)

# 定义模型
class RGCN(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_relations):
        super().__init__()
        self.conv1 = RGCNConv(
            in_channels=in_channels,
            out_channels=hidden_channels,
            num_relations=num_relations,
        )
        self.conv2 = RGCNConv(
            in_channels=hidden_channels,
            out_channels=out_channels,
            num_relations=num_relations
        )
        self.relu = nn.ReLU()
        
    def forward(self, x, edge_index, edge_type):
        x = self.relu(self.conv1(x, edge_index, edge_type))
        x = self.conv2(x, edge_index, edge_type)
        return x
    
    
# 前向传播
model = RGCN(3, 16, 8, 2)
out = model(data.x, data.edge_index, data.edge_type)
print(out)
深度学习 > PyTorch
#人工智能 #神经网络 #深度学习 #PyTorch #图神经网络 #PyG
PyG
https://blog.shinebook.net/2025/05/02/人工智能/pytorch/PyG/
作者
X
发布于
2025年5月2日
许可协议