PyTorch Geometric中InMemoryDataset数据加载问题解析

问题背景

在使用PyTorch Geometric（简称PyG）进行图神经网络开发时，开发者经常会遇到自定义数据集加载的问题。特别是当继承InMemoryDataset类创建自定义数据集时，一个常见但容易被忽视的问题是：预处理后的数据文件虽然存在，但数据集实例却显示为空。

问题现象

开发者定义了一个MolecularShiftsDataset类继承自InMemoryDataset，用于处理分子图数据。虽然预处理过程成功完成并生成了.pt文件，但当实例化数据集类时，数据集显示为空（长度为1且没有节点数据）。然而，当直接使用InMemoryDataset的load()方法加载同一文件时，数据却能正常加载。

根本原因

PyG的InMemoryDataset类设计上不会在__init__方法中自动加载预处理数据。这是有意为之的设计选择，目的是给予开发者更大的灵活性。开发者需要显式调用self.load()方法来加载预处理数据。

解决方案

对于自定义数据集类，正确的做法是在__init__方法中显式调用load()方法：

class MolecularShiftsDataset(InMemoryDataset):
    def __init__(
        self,
        dataset: pd.DataFrame,
        encoding: str,
        sdf: SDMolSupplier,
        name: str = "NMR_FF",
        root="data/graph_datasets/",
        edge_attr: bool = True,
        transform=None,
        pre_transform=None,
    ):
        self.encoding = encoding
        self.name = name
        self.dataset = dataset
        self.sdf = sdf
        self.edge_attr = edge_attr
        super().__init__(root, transform, pre_transform)
        self.load(self.processed_paths[0])  # 显式加载预处理数据

最佳实践

1. 明确加载时机：在__init__中立即加载数据，或者在首次访问数据时惰性加载
2. 错误处理：添加对预处理文件存在的检查
3. 文档注释：在自定义数据集类中明确说明数据加载行为

深入理解

PyG的这种设计有几个优点：

• 允许数据集类在不需要立即加载所有数据的情况下初始化
• 支持更复杂的数据加载策略
• 为大型数据集提供了灵活性，可以按需加载

对于分子图数据处理这种特定场景，建议在数据预处理阶段确保：

1. 分子结构到图的转换正确性
2. 特征编码的完整性
3. 数据分割的合理性

总结

PyTorch Geometric的InMemoryDataset提供了强大的自定义数据集支持，但需要开发者理解其设计理念。显式调用load()方法是正确处理预处理数据的关键。这种设计虽然增加了一点复杂性，但为各种图数据处理场景提供了更大的灵活性。