
PyTorch Geometric 从 CSV 构建异构图MovieLens 数据集加载全流程实战【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric导读在真实业务场景中图数据往往散落在多张关系型表格如 CSV中一张表存节点属性、另一张表存节点间关系。本文基于 PyTorch GeometricPyG官方教程 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教程内嵌代码与该脚本完全一致读者可以直接运行脚本复现本文全部步骤。二、数据准备下载并解压 MovieLensPyG 在 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指向的文件到folderextract_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 行movieIdtitlegenres1Toy Story (1995)Adventure|Animation|Children|Comedy|Fantasy2Jumanji (1995)Adventure|Children|Fantasy3Grumpier Old Men (1995)Comedy|Romance4Waiting to Exhale (1995)Comedy|Drama|Romance5Father of the Bride Part II (1995)ComedymovieId每部电影的唯一标识title电影标题可编码为文本特征genres电影类型用|分隔多个标签是多标签分类问题的天然素材。ratings.csv头部前 5 行userIdmovieIdratingtimestamp114.0964982703134.0964981247164.09649822241475.09649838151505.0964982931userId与movieId共同定义了谁给哪部电影打了分即边的两端rating是 0.55.0 的浮点评分可作为边特征甚至回归/分类的监督标签timestamp在本教程中因简化而被忽略但它保留了评分的时序信息在有时间建模需求时可另行利用。四、核心工具一load_node_csv—— 节点加载与索引映射图神经网络要求数据表示尽量紧凑节点编号必须是0 ~ num_rows-1的连续整数。为此定义load_node_csvimport torch def load_node_csv(path, index_col, encodersNone, **kwargs): df pd.read_csv(path, index_colindex_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工作原理解析pd.read_csv(path, index_colindex_col)将index_col指定的列设为 DataFrame 索引例如以movieId为索引mapping把原始 ID如movieId1映射为连续整数{原始ID: 行号}。这样第 0 行的电影特征恰好可以通过x[0]访问保证特征矩阵与edge_index中的编号一一对应encoders是一个{列名: 编码器对象}字典每个编码器把该列的原始值编码成[num_rows, dim_i]的张量最后沿最后一维torch.cat拼接成x形状为[num_nodes, num_features]。**kwargs会透传给pd.read_csv如sep、usecols等方便处理非标准 CSV 格式。五、核心工具二三种列编码器编码器Encoder是本节把任意列转成数值特征的统一抽象它接收一个 pandas Seriesdf[col]返回一个张量。5.1SequenceEncoder文本 → 预训练句向量电影标题是自然语言可借助sentence-transformers库加载预训练 NLP 模型编码为低维稠密向量。先安装依赖pip install sentence-transformersclass SequenceEncoder: def __init__(self, model_nameall-MiniLM-L6-v2, deviceNone): self.device device self.model SentenceTransformer(model_name, devicedevice) torch.no_grad() def __call__(self, df): x self.model.encode(df.values, show_progress_barTrue, convert_to_tensorTrue, deviceself.device) return x.cpu()默认模型all-MiniLM-L6-v2是 SBERT 生态中体积小、速度快、效果均衡的通用句嵌入模型输出 384 维向量整个编码过程包裹在torch.no_grad()下不参与梯度计算特征在构图阶段一次性固化convert_to_tensorTrue让输出直接是 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, dtypeNone): 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主体逻辑。六、生成节点数据并初始化HeteroData6.1 生成 movie 节点特征与 user 节点映射movie_x, movie_mapping load_node_csv( movie_path, index_colmovieId, encoders{ title: SequenceEncoder(), genres: GenresEncoder() }) _, user_mapping load_node_csv(rating_path, index_coluserId)电影的title用句向量编码、genres用 one-hot 编码拼接后得到movie_x形状[9742, 404]其中 384 维来自 SBERT 句向量、20 维来自类型 one-hot用户没有额外属性因此不传encodersx为None仅取回user_mapping。6.2 填充HeteroDatafrom 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_nodes610 }, 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_indexdef load_edge_csv(path, src_index_col, src_mapping, dst_index_col, dst_mapping, encodersNone, **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_coluserId, src_mappinguser_mapping, dst_index_colmovieId, dst_mappingmovie_mapping, encoders{rating: IdentityEncoder(dtypetorch.long)}, ) data[user, rates, movie].edge_index edge_index data[user, rates, movie].edge_label edge_label print(data)输出HeteroData( user{ num_nodes610 }, 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链接预测需要训练/验证/测试三个互不重叠的边集合。RandomLinkSplittorch_geometric/transforms/random_link_split.py实现边级随机划分from torch_geometric.transforms import RandomLinkSplit transform RandomLinkSplit( num_val0.05, num_test0.1, neg_sampling_ratio0.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.2neg_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_urlextract_zip原始 CSV 文件节点建模load_node_csv 自定义编码器节点特征x ID 映射表边建模load_edge_csv 自定义编码器edge_index 边特征图对象HeteroData字典式赋值可直接训练的异构图训练预处理ToUndirectedRandomLinkSplittrain/val/test 三个子图实践要点速记所有原始 ID 必须经mapping转为连续整数edge_index才能正确索引特征矩阵编码器接口统一为__call__(df) - Tensor文本、多标签、数值列都能以相同方式接入HeteroData的字符串键节点与三元组键边赋值语法是其嵌套字典设计的直接体现训练前的反向边补充与边级切分是链接预测任务防泄漏的两个关键步骤。将此流水线中的load_node_csv/load_edge_csv与各类编码器稍作修改即可迁移到电商用户-商品、社交网络用户-内容、生物信息化合物-靶点等任意多表关联 关系预测场景构建自己的异构图数据集。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考