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', name='Cora')

虽然PyG的Planetoid类提供了便捷的数据加载接口，但当我们面临以下场景时，理解原始数据处理流程就变得至关重要：

• 自定义数据预处理：官方预处理可能不符合我们的需求
• 调试与验证：当模型表现异常时，需要检查数据生成环节
• 迁移到新数据集：处理类似结构的非标准数据集
• 性能优化：控制数据加载的每个环节以提升效率

本文将带您从原始.tgz文件开始，逐步构建与PyG兼容的Data对象，涵盖文件下载、解析、邻接矩阵构建、特征/标签加载、掩码生成等完整流程。

2. 环境准备与数据下载

2.1 安装必要依赖

pip install torch torch-geometric numpy scipy requests tqdm

2.2 数据集下载链接

三个数据集的官方源文件可从LINQS网站获取：

2.3 自动化下载脚本

import os import requests import tarfile from tqdm import tqdm def download_dataset(url, save_path): """下载并解压数据集""" os.makedirs(save_path, exist_ok=True) filename = os.path.join(save_path, url.split('/')[-1]) # 下载文件 response = requests.get(url, stream=True) total_size = int(response.headers.get('content-length', 0)) with open(filename, 'wb') as f, tqdm( desc=filename, total=total_size, unit='iB', unit_scale=True, unit_divisor=1024, ) as bar: for data in response.iter_content(chunk_size=1024): size = f.write(data) bar.update(size) # 解压文件 with tarfile.open(filename, 'r:gz') as tar: tar.extractall(path=save_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 └── README

3.1 关键文件说明

1. .content文件（如cora.content）

   • 每行格式：<paper_id> <word_attributes>+ <class_label>
   • 示例：31336 0 0 ... 1 0 Neural_Networks

2. .cites文件（如cora.cites）

   • 每行格式：<cited_paper_id> <citing_paper_id>
   • 示例：35 1033表示论文1033引用了论文35

3.2 数据统计对比

4. 数据加载与预处理

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], dtype=np.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], dtype=np.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, dtype=np.int64).T # 转换为无向图 edge_index = np.concatenate([edge_index, edge_index[[1,0]]], axis=1) 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_ratio=0.05, val_ratio=0.18): """创建PyG Data对象并生成分割掩码""" # 转换为torch tensor x = torch.tensor(features, dtype=torch.float32) edge_index = torch.tensor(edge_index, dtype=torch.long) y = torch.tensor(y, dtype=torch.long) # 创建随机分割 num_nodes = y.size(0) num_classes = y.max().item() + 1 # 确保每个类至少有少量训练样本 train_mask = torch.zeros(num_nodes, dtype=torch.bool) val_mask = torch.zeros(num_nodes, dtype=torch.bool) test_mask = torch.zeros(num_nodes, dtype=torch.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_size+val_size]] = True test_mask[idx[train_size+val_size:]] = True # 构建Data对象 data = Data( x=x, edge_index=edge_index, y=y, train_mask=train_mask, val_mask=val_mask, test_mask=test_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', name='Cora')[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], dtype=np.int64) paper_to_idx = {paper_id: idx for idx, paper_id in enumerate(paper_ids)} return features, y, paper_to_idx, unique_labels

6.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, dtype=np.int64).T edge_index = np.concatenate([edge_index, edge_index[[1,0]]], axis=1) return features, edge_index, y, class_names

7. 完整项目结构建议

对于实际项目，建议采用如下目录结构：

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, name='cora', root='./data', force_reprocess=False): 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_ok=True) 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(f'Unknown 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