如何检测Python对象的内存占用？——针对含PyTorch张量的列表场景

嘿，我懂你为啥觉得sys.getsizeof()没用——它只能拿到Python对象本身那点“壳子”的大小，比如列表只是存了一堆元素的引用，每个引用在64位系统里才8字节，完全碰不到PyTorch张量（尤其是稀疏张量）实际存储数据的内存。下面给你俩靠谱的办法，能准确算出整个列表的总内存占用：

方法一：递归遍历，精准计算每个元素的内存 #

这个方法会递归遍历你的列表、元组，逐个计算每个张量（不管密集还是COO稀疏）的实际内存，最后累加总和。核心思路是：

• 对于密集张量：用numel()拿到总元素数，乘以element_size()（每个元素的字节数），就是它的实际内存。
• 对于COO稀疏张量：要单独计算indices和values两部分的内存，因为稀疏张量的核心数据都存在这两个组件里。

给你写好现成的函数，直接用就行：

import torch
import sys

def calculate_total_memory(obj):
    total_bytes = 0
    # 处理列表：递归遍历每个元素
    if isinstance(obj, list):
        for item in obj:
            total_bytes += calculate_total_memory(item)
    # 处理元组：同样递归遍历每个元素
    elif isinstance(obj, tuple):
        for elem in obj:
            total_bytes += calculate_total_memory(elem)
    # 处理PyTorch张量
    elif isinstance(obj, torch.Tensor):
        if obj.is_sparse:
            # COO稀疏张量：计算indices和values的内存
            total_bytes += obj.indices().numel() * obj.indices().element_size()
            total_bytes += obj.values().numel() * obj.values().element_size()
        else:
            # 密集张量：计算所有元素的总字节数
            total_bytes += obj.numel() * obj.element_size()
    # 其他Python基础对象，用sys.getsizeof兜底
    else:
        total_bytes += sys.getsizeof(obj)
    return total_bytes

# 举个测试例子
if __name__ == "__main__":
    # 创建测试用的张量
    dense_tensor1 = torch.randn(200, 200)  # 密集张量
    dense_tensor2 = torch.randn(100, 100)
    sparse_tensor = torch.randn(1000, 1000).to_sparse_coo()  # COO稀疏张量
    
    # 构造你的目标列表
    my_list = [(dense_tensor1, dense_tensor2, sparse_tensor)] * 3
    
    # 计算总内存
    total = calculate_total_memory(my_list)
    print(f"整个列表的内存占用：{total / (1024**2):.2f} MB")  # 转成MB更直观

方法二：利用PyTorch的内存统计（GPU场景适用） #

如果你的张量都在GPU上，可以用PyTorch自带的内存统计工具，通过创建列表前后的内存差值来计算占用。这个方法简单，但要注意测试时尽量隔离其他GPU操作，避免干扰结果：

import torch

# 先记录当前GPU已分配内存
before = torch.cuda.memory_allocated()

# 创建你的列表（这里用上面的测试数据示例）
dense_tensor1 = torch.randn(200, 200).cuda()
dense_tensor2 = torch.randn(100, 100).cuda()
sparse_tensor = torch.randn(1000, 1000).to_sparse_coo().cuda()
my_list = [(dense_tensor1, dense_tensor2, sparse_tensor)] * 3

# 计算差值
after = torch.cuda.memory_allocated()
occupied_memory = after - before

print(f"列表占用GPU内存：{occupied_memory / (1024**2):.2f} MB")

补充说明：为啥sys.getsizeof()不好用？ #

sys.getsizeof()只能返回对象在Python解释器堆中的大小，比如列表本身只是存储了元素的引用地址，每个地址占8字节（64位系统），但PyTorch张量的实际数据是存在PyTorch独立管理的内存池里（CPU或GPU），这部分内存完全不在Python堆的统计范围内，所以用它测出来的结果会严重偏小，完全没参考价值。

备注：内容来源于stack exchange，提问作者roz