如何排查PyTorch Geometric数据对象各字段的内存占用

先搞懂为什么删字段没差值 #

PyTorch Geometric的Data对象里，data.x、data.edge_index这类张量是由PyTorch内存管理的，但你添加的字符串属于Python原生对象，PyTorch的内存统计工具（包括get_data_size）完全不会计算这些。get_data_size说的“理论内存”，就是只算张量的字节数——比如data.x的元素数×每个元素的字节数，完全忽略Python层面的字符串、字典等对象。所以你删除字符串字段后，PyTorch的内存统计看不出变化，差值自然为0。

统计字符串内存的实用方法 #

1. 用Python原生工具直接计算 #

Python的sys.getsizeof能直接返回单个对象的内存字节数，针对字符串字段用这个就行：

import sys
import torch
from torch_geometric.data import Data

# 示例Data对象
data = Data(x=torch.randn(100, 16), edge_index=torch.randint(0, 100, (2, 200)))
data.db_path = "/user/data/datasets/graph_db_v2"
data.model_tag = "gcn_baseline"

# 单独统计每个字符串字段
print(f"db_path 内存: {sys.getsizeof(data.db_path)} bytes")
print(f"model_tag 内存: {sys.getsizeof(data.model_tag)} bytes")

注：sys.getsizeof只计算对象本身的内存，字符串是不可变对象，这个结果足够评估单个字段的消耗。

2. 批量扫描所有非张量字段 #

如果添加了多个附加字段，不想逐个手动统计，就遍历Data对象的属性，跳过张量和私有属性批量计算：

import sys
import torch
from torch_geometric.data import Data

data = Data(x=torch.randn(100, 16), edge_index=torch.randint(0, 100, (2, 200)))
data.db_path = "/user/data/datasets/graph_db_v2"
data.model_tag = "gcn_baseline"
data.train_note = "lr=0.01, batch_size=64"

for attr in dir(data):
    # 跳过下划线开头的私有属性，以及张量类型的属性
    if attr.startswith('_') or isinstance(getattr(data, attr), torch.Tensor):
        continue
    val = getattr(data, attr)
    print(f"字段 {attr}: {sys.getsizeof(val)} bytes")

3. 精简字符串内存的小技巧 #

要是发现字符串占用的内存影响批大小，试试这些优化：

• 重复字符串共用引用：多个Data对象用同一个模型名称时，直接赋值同一个变量，Python会自动复用内存
• 用短编码替代长字符串：把模型名、路径映射成整数ID，训练时只存ID，需要显示时再反向映射
• 删除训练无用字段：比如数据库路径，数据加载完成后直接执行del data.db_path，避免无用内存占用

内容的提问来源于stack exchange，提问作者Knowledge seeker