如何为PyTorch Geometric的ZINC数据集Data对象添加新特征

要将新特征列表整合到ZINC数据集的每个Data对象中，可以尝试以下两种可靠方法：

方法一：原地修改数据集元素 #

遍历数据集的每个元素，直接为其添加新特征属性，注意要保证新特征张量的设备（CPU/GPU）与Data对象的其他张量一致：

import torch
from torch_geometric.datasets import ZINC

# 加载原始数据集
zinc_dataset = ZINC(root='my_path', split='train')

# 假设new_features_list是你预先计算好的新特征列表，每个元素对应一个图的节点特征张量
new_features_list = [torch.randn(data.x.shape[0], 12) for data in zinc_dataset]  # 示例数据

# 遍历添加新特征
for data, new_feat in zip(zinc_dataset, new_features_list):
    # 同步张量设备
    new_feat = new_feat.to(data.x.device)
    # 为Data对象添加新属性
    data.new_feature = new_feat

# 验证修改结果
print(zinc_dataset[0])

方法二：创建新的InMemoryDataset（推荐） #

如果原地修改后未生效，可能是因为ZINC属于InMemoryDataset，内部数据存储在_data_list中，直接修改外部引用可能无法同步到内部存储。可以通过重建数据集确保修改被正确保存：

from torch_geometric.data import InMemoryDataset, Data
from torch_geometric.datasets import ZINC

# 加载原始数据集
zinc_dataset = ZINC(root='my_path', split='train')
new_features_list = [torch.randn(data.x.shape[0], 12) for data in zinc_dataset]  # 示例数据

# 构造包含新特征的Data对象列表
new_data_list = []
for data, new_feat in zip(zinc_dataset, new_features_list):
    new_feat = new_feat.to(data.x.device)
    # 合并原有属性与新特征，创建新的Data对象
    new_data = Data(
        x=data.x,
        edge_index=data.edge_index,
        edge_attr=data.edge_attr,
        y=data.y,
        new_feature=new_feat
    )
    new_data_list.append(new_data)

# 自定义修改后的数据集类
class ModifiedZINC(InMemoryDataset):
    def __init__(self, root, data_list):
        super().__init__(root)
        self.data, self.slices = self.collate(data_list)

# 初始化新数据集
modified_zinc = ModifiedZINC(root='my_path_modified', data_list=new_data_list)

# 验证结果
print(modified_zinc[0])

常见问题排查 #

• 确保新特征的节点维度与对应图的节点数匹配（即new_feature的第一维度等于data.x的第一维度）。
• 检查张量设备是否一致：若数据集已移至GPU，新特征需同步到相同设备，否则会引发报错。
• 之前的方案无效，大概率是未处理设备一致性，或未针对InMemoryDataset的内部存储机制修改，上述方法可规避这类问题。

内容的提问来源于stack exchange，提问作者Alexandre Bloch