Cora/CiteSeer/PubMed 原始数据处理：从 .tgz 到 PyG Data 对象的 5 步转换指南

Cora/CiteSeer/PubMed 原始数据到 PyG Data 对象的完整转换指南1. 引言为什么需要手动处理原始数据在大多数图神经网络GNN教程中我们通常会看到这样简单的数据加载代码from torch_geometric.datasets import Planetoid dataset Planetoid(root/tmp/Cora, nameCora)虽然PyG的Planetoid类提供了便捷的数据加载接口但当我们面临以下场景时理解原始数据处理流程就变得至关重要自定义数据预处理官方预处理可能不符合我们的需求调试与验证当模型表现异常时需要检查数据生成环节迁移到新数据集处理类似结构的非标准数据集性能优化控制数据加载的每个环节以提升效率本文将带您从原始.tgz文件开始逐步构建与PyG兼容的Data对象涵盖文件下载、解析、邻接矩阵构建、特征/标签加载、掩码生成等完整流程。2. 环境准备与数据下载2.1 安装必要依赖pip install torch torch-geometric numpy scipy requests tqdm2.2 数据集下载链接三个数据集的官方源文件可从LINQS网站获取数据集下载链接Corahttps://linqs-data.soe.ucsc.edu/public/lbc/cora.tgzCiteSeerhttps://linqs-data.soe.ucsc.edu/public/lbc/citeseer.tgzPubMedhttps://linqs-data.soe.ucsc.edu/public/Pubmed-Diabetes.tgz2.3 自动化下载脚本import os import requests import tarfile from tqdm import tqdm def download_dataset(url, save_path): 下载并解压数据集 os.makedirs(save_path, exist_okTrue) filename os.path.join(save_path, url.split(/)[-1]) # 下载文件 response requests.get(url, streamTrue) total_size int(response.headers.get(content-length, 0)) with open(filename, wb) as f, tqdm( descfilename, totaltotal_size, unitiB, unit_scaleTrue, unit_divisor1024, ) as bar: for data in response.iter_content(chunk_size1024): size f.write(data) bar.update(size) # 解压文件 with tarfile.open(filename, r:gz) as tar: tar.extractall(pathsave_path) os.remove(filename) # 删除压缩包 # 示例下载Cora数据集 download_dataset( https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz, ./raw_data/cora )3. 文件结构解析解压后的典型文件结构如下以Cora为例cora/ ├── cora.cites ├── cora.content └── README3.1 关键文件说明.content文件如cora.content每行格式paper_id word_attributes class_label示例31336 0 0 ... 1 0 Neural_Networks.cites文件如cora.cites每行格式cited_paper_id citing_paper_id示例35 1033表示论文1033引用了论文353.2 数据统计对比指标CoraCiteSeerPubMed节点数2,7083,32719,717边数5,4294,73244,338特征维度1,4333,703500类别数7634. 数据加载与预处理4.1 加载.content文件import numpy as np from scipy.sparse import csr_matrix def load_content(file_path): 加载.content文件并返回特征矩阵和标签 with open(file_path, r) as f: lines f.read().splitlines() # 解析每行数据 data [line.split(\t) for line in lines] paper_ids [item[0] for item in data] # 构建特征矩阵 features np.array([list(map(int, item[1:-1])) for item in data], dtypenp.float32) # 处理标签 labels [item[-1] for item in data] unique_labels sorted(list(set(labels))) label_to_idx {label: idx for idx, label in enumerate(unique_labels)} y np.array([label_to_idx[label] for label in labels], dtypenp.int64) # 创建paper_id到索引的映射 paper_to_idx {paper_id: idx for idx, paper_id in enumerate(paper_ids)} return features, y, paper_to_idx, unique_labels # 示例使用 features, y, paper_to_idx, class_names load_content(./raw_data/cora/cora.content)4.2 加载.cites文件构建邻接矩阵def load_cites(file_path, paper_to_idx): 加载.cites文件并构建邻接矩阵 with open(file_path, r) as f: lines f.read().splitlines() edges [] for line in lines: cited, citing line.split(\t) if cited in paper_to_idx and citing in paper_to_idx: edges.append([paper_to_idx[citing], paper_to_idx[cited]]) # 注意方向 edge_index np.array(edges, dtypenp.int64).T # 转换为无向图 edge_index np.concatenate([edge_index, edge_index[[1,0]]], axis1) return edge_index # 示例使用 edge_index load_cites(./raw_data/cora/cora.cites, paper_to_idx)5. 构建PyG Data对象5.1 完整转换代码import torch from torch_geometric.data import Data def create_pyg_data(features, edge_index, y, train_ratio0.05, val_ratio0.18): 创建PyG Data对象并生成分割掩码 # 转换为torch tensor x torch.tensor(features, dtypetorch.float32) edge_index torch.tensor(edge_index, dtypetorch.long) y torch.tensor(y, dtypetorch.long) # 创建随机分割 num_nodes y.size(0) num_classes y.max().item() 1 # 确保每个类至少有少量训练样本 train_mask torch.zeros(num_nodes, dtypetorch.bool) val_mask torch.zeros(num_nodes, dtypetorch.bool) test_mask torch.zeros(num_nodes, dtypetorch.bool) for c in range(num_classes): idx (y c).nonzero().view(-1) idx idx[torch.randperm(idx.size(0))] train_size max(1, int(train_ratio * idx.size(0))) val_size int(val_ratio * idx.size(0)) train_mask[idx[:train_size]] True val_mask[idx[train_size:train_sizeval_size]] True test_mask[idx[train_sizeval_size:]] True # 构建Data对象 data Data( xx, edge_indexedge_index, yy, train_masktrain_mask, val_maskval_mask, test_masktest_mask ) return data # 完整流程示例 features, y, paper_to_idx, _ load_content(./raw_data/cora/cora.content) edge_index load_cites(./raw_data/cora/cora.cites, paper_to_idx) cora_data create_pyg_data(features, edge_index, y)5.2 与官方数据对比验证from torch_geometric.datasets import Planetoid # 加载官方数据 official_data Planetoid(root/tmp/Cora, nameCora)[0] # 对比关键属性 print(官方数据 vs 手动处理数据对比:) print(f节点数: {official_data.num_nodes} {cora_data.num_nodes}) print(f边数: {official_data.num_edges} ≈ {cora_data.num_edges} (方向处理可能导致差异)) print(f特征维度: {official_data.num_features} {cora_data.num_features}) print(f类别数: {official_data.num_classes} {len(class_names)})6. 高级处理技巧6.1 处理稀疏特征矩阵对于大型数据集如PubMed特征矩阵可能非常稀疏使用稀疏矩阵存储可节省内存from scipy.sparse import csr_matrix def load_content_sparse(file_path): 稀疏矩阵方式加载.content文件 with open(file_path, r) as f: lines f.read().splitlines() data [line.split(\t) for line in lines] paper_ids [item[0] for item in data] # 构建稀疏矩阵 rows, cols, values [], [], [] for row_idx, item in enumerate(data): for col_idx, val in enumerate(item[1:-1]): if int(val) 1: rows.append(row_idx) cols.append(col_idx) values.append(1.0) num_nodes len(data) num_features len(data[0]) - 2 # 减去paper_id和label features csr_matrix((values, (rows, cols)), shape(num_nodes, num_features)) # 处理标签 labels [item[-1] for item in data] unique_labels sorted(list(set(labels))) label_to_idx {label: idx for idx, label in enumerate(unique_labels)} y np.array([label_to_idx[label] for label in labels], dtypenp.int64) paper_to_idx {paper_id: idx for idx, paper_id in enumerate(paper_ids)} return features, y, paper_to_idx, unique_labels6.2 处理CiteSeer的特殊情况CiteSeer数据集中存在一些孤立节点需要特殊处理def process_citeseer(content_path, cites_path): 处理CiteSeer的特殊情况 features, y, paper_to_idx, class_names load_content(content_path) # 加载引用关系 with open(cites_path, r) as f: lines f.read().splitlines() # 找出有连接的节点 connected_nodes set() for line in lines: cited, citing line.split(\t) if cited in paper_to_idx and citing in paper_to_idx: connected_nodes.add(paper_to_idx[cited]) connected_nodes.add(paper_to_idx[citing]) # 创建新索引 connected_nodes sorted(connected_nodes) new_idx_map {old_idx: new_idx for new_idx, old_idx in enumerate(connected_nodes)} # 过滤特征和标签 features features[connected_nodes] y y[connected_nodes] # 重建边索引 edges [] for line in lines: cited, citing line.split(\t) if cited in paper_to_idx and citing in paper_to_idx: edges.append([new_idx_map[paper_to_idx[citing]], new_idx_map[paper_to_idx[cited]]]) edge_index np.array(edges, dtypenp.int64).T edge_index np.concatenate([edge_index, edge_index[[1,0]]], axis1) return features, edge_index, y, class_names7. 完整项目结构建议对于实际项目建议采用如下目录结构graph_data_pipeline/ ├── data/ # 原始数据 │ ├── cora/ │ ├── citeseer/ │ └── pubmed/ ├── processed/ # 处理后的数据 ├── utils/ │ ├── __init__.py │ ├── data_loader.py # 数据加载工具 │ └── preprocess.py # 预处理工具 ├── configs/ # 配置文件 │ └── data_config.yaml ├── tests/ # 单元测试 ├── main.py # 主程序 └── requirements.txt在data_loader.py中可以实现一个统一的数据加载接口class GraphDataset: def __init__(self, namecora, root./data, force_reprocessFalse): self.name name.lower() self.raw_dir os.path.join(root, self.name) self.processed_dir os.path.join(root, processed) os.makedirs(self.processed_dir, exist_okTrue) self.processed_file os.path.join( self.processed_dir, f{self.name}_data.pt ) if not os.path.exists(self.processed_file) or force_reprocess: self.process() else: self.data torch.load(self.processed_file) def process(self): if self.name cora: features, edge_index, y, _ process_cora( os.path.join(self.raw_dir, cora.content), os.path.join(self.raw_dir, cora.cites) ) elif self.name citeseer: features, edge_index, y, _ process_citeseer( os.path.join(self.raw_dir, citeseer.content), os.path.join(self.raw_dir, citeseer.cites) ) elif self.name pubmed: features, edge_index, y, _ process_pubmed( os.path.join(self.raw_dir, Pubmed-Diabetes.NODE.paper.tab), os.path.join(self.raw_dir, Pubmed-Diabetes.DIRECTED.cites.tab) ) else: raise ValueError(fUnknown dataset: {self.name}) self.data create_pyg_data(features, edge_index, y) torch.save(self.data, self.processed_file) def get(self): return self.data