PyTorch Geometric 从 CSV 构建异构图：MovieLens 数据集加载全流程实战

PyTorch Geometric 从 CSV 构建异构图：MovieLens 数据集加载全流程实战

【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric

导读

在真实业务场景中，图数据往往散落在多张关系型表格（如 CSV）中：一张表存节点属性、另一张表存节点间关系。本文基于 PyTorch Geometric（PyG）官方教程 docs/source/tutorial/load_csv.rst 及其配套示例脚本 examples/hetero/load_csv.py，完整演示如何把 MovieLens 的movies.csv与ratings.csv两张表转换为 PyG 的HeteroData异构图对象。读完本文，你将掌握：CSV 索引列到连续整数 ID 的映射方法、可插拔的列编码器（文本嵌入、多标签 one-hot、原始数值直通）设计、HeteroData节点/边类型的赋值语法，以及为后续链接预测训练做数据切分的完整思路。

---

一、任务背景：从表格数据到异构图

MovieLens 数据集（由 GroupLens 研究组维护）是一个经典的推荐系统评测数据。本文使用的ml-latest-small子集包含约100k 条评分记录，覆盖9000+ 部电影和600+ 用户。我们要完成如下图建模：

• 两种节点类型：user（用户）、movie（电影）；
• 一种边类型：(user, rates, movie)，表示"某个用户给某部电影打了分"。

这就是一个典型的异构图（Heterogeneous Graph）——节点和边携带不同类型，可进一步作为异构图表征模型（如 RGCN、HGT 等）的输入。

注：原文档中提到的配套可执行脚本位于 examples/hetero/load_csv.py，教程内嵌代码与该脚本完全一致，读者可以直接运行脚本复现本文全部步骤。

---

二、数据准备：下载并解压 MovieLens

PyG 在 torch_geometric/data/download.py 与 torch_geometric/data/extract.py 中提供了download_url与extract_zip两个实用工具。前者基于urllib分块（10MB/次）下载并支持断点复用（文件已存在时直接返回路径，见 download.py），后者将 zip 归档解压到指定目录。

from torch_geometric.data import download_url, extract_zip url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip' extract_zip(download_url(url, '.'), '.') movie_path = './ml-latest-small/movies.csv' rating_path = './ml-latest-small/ratings.csv'

在示例脚本中，数据集被下载到仓库外的data/MovieLens目录：

root = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/MovieLens') extract_zip(download_url(url, root), root)

参数说明：download_url(url, folder)下载url指向的文件到folder；extract_zip(path, folder)将 zip 解压到folder。返回的 CSV 路径随下载目录变化，后续代码只需关注movie_path/rating_path两个变量。

---

三、审视原始数据：两张 CSV 的结构

先用 pandas 快速查看两个文件的前几行，确认列结构：

import pandas as pd print(pd.read_csv(movie_path).head()) print(pd.read_csv(rating_path).head())

movies.csv头部（前 5 行）

movieId	title	genres
1	Toy Story (1995)	Adventure|Animation|Children|Comedy|Fantasy
2	Jumanji (1995)	Adventure|Children|Fantasy
3	Grumpier Old Men (1995)	Comedy|Romance
4	Waiting to Exhale (1995)	Comedy|Drama|Romance
5	Father of the Bride Part II (1995)	Comedy

• movieId：每部电影的唯一标识；
• title：电影标题（可编码为文本特征）；
• genres：电影类型，用|分隔多个标签（是多标签分类问题的天然素材）。

ratings.csv头部（前 5 行）

userId	movieId	rating	timestamp
1	1	4.0	964982703
1	3	4.0	964981247
1	6	4.0	964982224
1	47	5.0	964983815
1	50	5.0	964982931

• userId与movieId共同定义了"谁给哪部电影打了分"，即边的两端；
• rating是 0.5～5.0 的浮点评分，可作为边特征（甚至回归/分类的监督标签）；
• timestamp在本教程中因简化而被忽略，但它保留了评分的时序信息，在有时间建模需求时可另行利用。

---

四、核心工具一：load_node_csv—— 节点加载与索引映射

图神经网络要求数据表示尽量紧凑：节点编号必须是0 ~ num_rows-1的连续整数。为此定义load_node_csv：

import torch def load_node_csv(path, index_col, encoders=None, **kwargs): df = pd.read_csv(path, index_col=index_col, **kwargs) mapping = {index: i for i, index in enumerate(df.index.unique())} x = None if encoders is not None: xs = [encoder(df[col]) for col, encoder in encoders.items()] x = torch.cat(xs, dim=-1) return x, mapping

工作原理解析：

1. pd.read_csv(path, index_col=index_col)将index_col指定的列设为 DataFrame 索引，例如以movieId为索引；
2. mapping把原始 ID（如movieId=1）映射为连续整数{原始ID: 行号}。这样第 0 行的电影特征恰好可以通过x[0]访问，保证特征矩阵与edge_index中的编号一一对应；
3. encoders是一个{列名: 编码器对象}字典，每个编码器把该列的原始值编码成[num_rows, dim_i]的张量，最后沿最后一维torch.cat拼接成x，形状为[num_nodes, num_features]。

**kwargs会透传给pd.read_csv（如sep、usecols等），方便处理非标准 CSV 格式。

---

五、核心工具二：三种列编码器

编码器（Encoder）是本节"把任意列转成数值特征"的统一抽象，它接收一个 pandas Series（df[col]），返回一个张量。

5.1SequenceEncoder：文本 → 预训练句向量

电影标题是自然语言，可借助sentence-transformers库加载预训练 NLP 模型编码为低维稠密向量。先安装依赖：

pip install sentence-transformers

class SequenceEncoder: def __init__(self, model_name='all-MiniLM-L6-v2', device=None): self.device = device self.model = SentenceTransformer(model_name, device=device) @torch.no_grad() def __call__(self, df): x = self.model.encode(df.values, show_progress_bar=True, convert_to_tensor=True, device=self.device) return x.cpu()

• 默认模型all-MiniLM-L6-v2是 SBERT 生态中体积小、速度快、效果均衡的通用句嵌入模型（输出 384 维向量）；
• 整个编码过程包裹在@torch.no_grad()下，不参与梯度计算，特征在构图阶段一次性固化；
• convert_to_tensor=True让输出直接是 PyTorch 张量，最后.cpu()移回内存，避免 GPU 显存长期占用。

替换model_name即可使用 sentence-transformers 支持的其他模型，但需注意特征维度会随之变化。

5.2GenresEncoder：多标签 → one-hot 特征

电影类型形如Adventure|Children|Fantasy，是可变数量的多标签。策略是先统计全数据集的类型全集，构造 one-hot 矩阵：

class GenresEncoder: def __init__(self, sep='|'): self.sep = sep def __call__(self, df): genres = set(g for col in df.values for g in col.split(self.sep)) mapping = {genre: i for i, genre in enumerate(genres)} x = torch.zeros(len(df), len(mapping)) for i, col in enumerate(df.values): for genre in col.split(self.sep): x[i, mapping[genre]] = 1 return x

输出形状为[num_movies, num_genres]，若电影i包含类型j则x[i, j] = 1。sep参数控制分隔符，可适配,、;等其他分隔格式。由于set迭代顺序不稳定，这里生成的mapping顺序每次运行可能不同，但特征语义一致，不影响模型训练。

5.3IdentityEncoder：原始数值 → 直通张量

评分本身是浮点数，无需变换，直接转成张量即可：

class IdentityEncoder: def __init__(self, dtype=None): self.dtype = dtype def __call__(self, df): return torch.from_numpy(df.values).view(-1, 1).to(self.dtype)

view(-1, 1)把一维 Series 变成列向量[N, 1]，dtype可指定为torch.float或torch.long（示例脚本中用torch.long将评分转为整数标签，便于后续分类任务）。

5.4 三种编码器的设计启示

从代码结构看（examples/hetero/load_csv.py），三个编码器共享同一接口——__call__(self, df) -> Tensor。这意味着任何满足该签名的自定义编码器都可即插即用：例如对数值列做 Min-Max 归一化、对类别列做 sklearn LabelEncoder、对图像列提取 CNN 特征等，无需改动load_node_csv/load_edge_csv主体逻辑。

---

六、生成节点数据并初始化HeteroData

6.1 生成 movie 节点特征与 user 节点映射

movie_x, movie_mapping = load_node_csv( movie_path, index_col='movieId', encoders={ 'title': SequenceEncoder(), 'genres': GenresEncoder() }) _, user_mapping = load_node_csv(rating_path, index_col='userId')

• 电影的title用句向量编码、genres用 one-hot 编码，拼接后得到movie_x（形状[9742, 404]，其中 384 维来自 SBERT 句向量、20 维来自类型 one-hot）；
• 用户没有额外属性，因此不传encoders，x为None，仅取回user_mapping。

6.2 填充HeteroData

from torch_geometric.data import HeteroData data = HeteroData() data['user'].num_nodes = len(user_mapping) # Users do not have any features. data['movie'].x = movie_x print(data)

输出：

HeteroData( user={ num_nodes=610 }, movie={ x[9742, 404] } )

关于HeteroData的实现细节，可参考 torch_geometric/data/hetero_data.py：它"试图模拟一个嵌套的 Python 字典"——字符串键对应节点存储（NodeStorage），三元组键对应边存储（EdgeStorage）。其__getitem__/__setitem__支持三类赋值方式（见 hetero_data.py）：

# (1) 初始化后逐个赋值 data = HeteroData() data['movie'].x = movie_x # (2) 关键字参数形式 data = HeteroData(movie={ 'x': movie_x }) # (3) 字典形式 data = HeteroData({'movie': { 'x': movie_x }})

由于user节点没有特征，只设置num_nodes。后续训练时，通常用torch.nn.Embedding为每个用户学习端到端的可训练嵌入（embedding），这正是异构推荐模型的常见做法。

---

七、核心工具三：load_edge_csv—— 边加载与索引重映射

ratings.csv中的userId/movieId是原始 ID，必须借助前面生成的user_mapping/movie_mapping重映射为连续编号，才能构造合法的edge_index：

def load_edge_csv(path, src_index_col, src_mapping, dst_index_col, dst_mapping, encoders=None, **kwargs): df = pd.read_csv(path, **kwargs) src = [src_mapping[index] for index in df[src_index_col]] dst = [dst_mapping[index] for index in df[dst_index_col]] edge_index = torch.tensor([src, dst]) edge_attr = None if encoders is not None: edge_attrs = [encoder(df[col]) for col, encoder in encoders.items()] edge_attr = torch.cat(edge_attrs, dim=-1) return edge_index, edge_attr

关键点：

• src_index_col/dst_index_col分别是源节点、目标节点的索引列名；
• src_mapping/dst_mapping是两个节点类型的 ID 映射表，逐行把原始 ID 翻译成连续索引；
• 返回的edge_index形状为[2, num_ratings]，第一行是源节点编号、第二行是目标节点编号；
• 与load_node_csv一致，encoders负责生成可选的边特征edge_attr。

7.1 组装评分边

edge_index, edge_label = load_edge_csv( rating_path, src_index_col='userId', src_mapping=user_mapping, dst_index_col='movieId', dst_mapping=movie_mapping, encoders={'rating': IdentityEncoder(dtype=torch.long)}, ) data['user', 'rates', 'movie'].edge_index = edge_index data['user', 'rates', 'movie'].edge_label = edge_label print(data)

输出：

HeteroData( user={ num_nodes=610 }, movie={ x=[9742, 404] }, (user, rates, movie)={ edge_index=[2, 100836], edge_label=[100836, 1] } )

这里data['user', 'rates', 'movie']的三元组语法定义了边类型：源节点类型user、关系名rates、目标节点类型movie。edge_label是评分张量（每行一个评分，形状[100836, 1]），既可当边特征，也可作为链接预测/评分预测的监督标签。

---

八、进一步：为图模型训练做准备

教程 docs/source/tutorial/load_csv.rst 生成的HeteroData已可直接输入异构模型；配套脚本 examples/hetero/load_csv.py 还演示了训练前必须的两步预处理：

8.1 补充反向边（ToUndirected）

多数 GNN 消息传递依赖双向邻接。ToUndirected变换（torch_geometric/transforms/to_undirected.py）会为每个边类型生成反向连接：对二部边(user, rates, movie)，自动创建(movie, rev_rates, user)：

from torch_geometric.transforms import ToUndirected data = ToUndirected()(data) del data['movie', 'rev_rates', 'user'].edge_label # 移除"反向"标签，防止标签泄漏

反向边的edge_label必须删除——评分语义只在user → movie方向有意义，反向边仅用于消息传递。

8.2 链接级数据切分（RandomLinkSplit）

链接预测需要训练/验证/测试三个互不重叠的边集合。RandomLinkSplit（torch_geometric/transforms/random_link_split.py）实现边级随机划分：

from torch_geometric.transforms import RandomLinkSplit transform = RandomLinkSplit( num_val=0.05, num_test=0.1, neg_sampling_ratio=0.0, edge_types=[('user', 'rates', 'movie')], rev_edge_types=[('movie', 'rev_rates', 'user')], ) train_data, val_data, test_data = transform(data)

参数说明（均可在源码 docstring 中确认）：

• num_val/num_test：验证/测试边比例，取值[0, 1]时按比例划分（默认分别为 0.1 / 0.2）；
• neg_sampling_ratio：负样本采样比例，设为0.0表示不生成负边（此处评分标签本身就是监督信号）；
• edge_types/rev_edge_types：指定参与切分的边类型及其反向边类型，确保反向边被同步划分，防止训练/测试间的数据泄漏；
• 其他可用参数：is_undirected（无向图防泄漏）、key（标签属性名，默认edge_label）、add_negative_train_samples、disjoint_train_ratio等。

切分后train_data/val_data/test_data各自携带对应子集的edge_index与edge_label，可直接送入异构图模型训练。

---

九、结语：一条可复用的"表格 → 异构图"流水线

回顾整个流程，load_csv.rst 教给我们的是一套高度通用的方法论，而不仅是 MovieLens 这一个例子：

环节	工具	关键产出
下载解压	download_url+extract_zip	原始 CSV 文件
节点建模	load_node_csv+ 自定义编码器	节点特征x+ ID 映射表
边建模	load_edge_csv+ 自定义编码器	edge_index+ 边特征
图对象	HeteroData字典式赋值	可直接训练的异构图
训练预处理	ToUndirected+RandomLinkSplit	train/val/test 三个子图

实践要点速记：

1. 所有原始 ID 必须经mapping转为连续整数，edge_index才能正确索引特征矩阵；
2. 编码器接口统一为__call__(df) -> Tensor，文本、多标签、数值列都能以相同方式接入；
3. HeteroData的字符串键（节点）与三元组键（边）赋值语法，是其"嵌套字典"设计的直接体现；
4. 训练前的反向边补充与边级切分是链接预测任务防泄漏的两个关键步骤。

将此流水线中的load_node_csv/load_edge_csv与各类编码器稍作修改，即可迁移到电商（用户-商品）、社交网络（用户-内容）、生物信息（化合物-靶点）等任意"多表关联 + 关系预测"场景，构建自己的异构图数据集。

【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric