【Bug已解决】Pytorch: Why is the memory occupied by the tensor variable so small? 解决方案-二趣网

【Bug已解决】Pytorch: Why is the memory occupied by the tensor variable so small? 解决方案

问题描述

在 PyTorch 中，开发者经常发现一个看似矛盾的现象：创建了一个很大的 Tensor（例如包含数百万个元素），但使用sys.getsizeof()检查时，发现 Tensor 对象本身只占用了很少的内存（通常只有几十字节）。这引发了一个常见疑问：

import torch import sys # 创建一个大型 Tensor big_tensor = torch.randn(1000000, 100) # 1亿个 float32 元素 print(f"Tensor 元素数量: {big_tensor.numel()}") print(f"Tensor 实际数据大小: {big_tensor.element_size() * big_tensor.numel() / 1e9:.2f} GB") print(f"sys.getsizeof(tensor): {sys.getsizeof(big_tensor)} bytes") # 输出可能只有几十字节！

为什么sys.getsizeof()报告的内存如此之小？Tensor 的实际数据存储在哪里？如何正确测量 Tensor 的内存占用？

错误复现

场景一：误用 sys.getsizeof

import torch import sys tensor = torch.randn(1000, 1000) print(f"sys.getsizeof: {sys.getsizeof(tensor)}") # 可能只有 72 bytes # 开发者误以为 Tensor 只占 72 字节，实际上数据占 1000*1000*4 = 4MB

场景二：误用sizeof

tensor = torch.randn(1000, 1000) print(f"__sizeof__: {tensor.__sizeof__()}") # 同样很小

场景三：内存监控不准确

import psutil import os process = psutil.Process(os.getpid()) mem_before = process.memory_info().rss tensor = torch.randn(10000, 10000) # ~400MB mem_after = process.memory_info().rss print(f"内存增量: {(mem_after - mem_before) / 1e6:.1f} MB") # 这个值可能不准确，因为操作系统内存分配策略

根因分析

1. Python 对象与底层数据分离

PyTorch Tensor 是一个 Python 对象，它包含：

元数据（形状、dtype、stride、device 等）：存储在 Python 堆中
实际数据：存储在独立的内存区域中（通过Storage对象管理）

sys.getsizeof()只测量 Python 对象本身的大小（即元数据部分），不包括底层Storage指向的实际数据。

2. Tensor 的内存结构

Tensor (Python 对象, ~72 bytes) ├── 元数据: shape, dtype, stride, device, requires_grad... └── Storage (数据存储对象) └── 实际数据 (element_size × numel bytes) └── 存储在 CPU RAM 或 GPU 显存中

3. Storage 的共享机制

多个 Tensor 可以共享同一个Storage（例如通过view、slice、transpose创建的 Tensor）。因此，Tensor 对象本身很小，实际数据由Storage管理。

4. CUDA Tensor 的特殊性

CUDA Tensor 的数据存储在 GPU 显存中，sys.getsizeof()完全无法反映 GPU 显存占用。

解决方案

方案一：使用 element_size() × numel() 计算实际大小

import torch def get_tensor_size_bytes(tensor): """获取 Tensor 实际数据大小（字节）""" return tensor.element_size() * tensor.nelement() tensor = torch.randn(1000, 1000) size = get_tensor_size_bytes(tensor) print(f"实际数据大小: {size} bytes = {size / 1e6:.2f} MB")

方案二：使用 storage().nbytes()

tensor = torch.randn(1000, 1000) print(f"Storage 大小: {tensor.storage().nbytes()} bytes") # 注意：storage().nbytes() 返回的是 Storage 的总大小 # 可能大于 Tensor 实际使用的部分（如有 padding）

方案三：使用 nbytes 属性（PyTorch 2.0+）

tensor = torch.randn(1000, 1000) print(f"nbytes: {tensor.nbytes} bytes") # PyTorch 2.0+

方案四：GPU 显存监控

import torch # 分配前 torch.cuda.reset_peak_memory_stats() mem_before = torch.cuda.memory_allocated() tensor = torch.randn(10000, 10000, device='cuda') # 分配后 mem_after = torch.cuda.memory_allocated() print(f"GPU 显存增量: {(mem_after - mem_before) / 1e6:.2f} MB") print(f"GPU 峰值显存: {torch.cuda.max_memory_allocated() / 1e6:.2f} MB")

完整修复代码

import torch import sys import os import psutil def get_tensor_memory(tensor): """ 全面获取 Tensor 的内存信息。 返回: dict: 包含各种内存度量 """ info = { 'python_object_size': sys.getsizeof(tensor), 'element_size': tensor.element_size(), 'numel': tensor.nelement(), 'data_size_bytes': tensor.element_size() * tensor.nelement(), 'data_size_mb': tensor.element_size() * tensor.nelement() / (1024**2), 'dtype': str(tensor.dtype), 'shape': tuple(tensor.shape), 'device': str(tensor.device), } # Storage 信息 if tensor.storage() is not None: info['storage_nbytes'] = tensor.storage().nbytes() else: info['storage_nbytes'] = 0 # nbytes 属性（PyTorch 2.0+） if hasattr(tensor, 'nbytes'): info['nbytes'] = tensor.nbytes # GPU 显存 if tensor.is_cuda: info['gpu_allocated'] = torch.cuda.memory_allocated(tensor.device) info['gpu_reserved'] = torch.cuda.memory_reserved(tensor.device) return info def print_tensor_memory(tensor, name="tensor"): """打印 Tensor 内存信息""" info = get_tensor_memory(tensor) print(f"\n--- {name} 内存信息 ---") print(f" Python 对象大小: {info['python_object_size']} bytes") print(f" 数据大小: {info['data_size_bytes']:,} bytes ({info['data_size_mb']:.2f} MB)") print(f" Storage 大小: {info['storage_nbytes']:,} bytes") print(f" 元素数量: {info['numel']:,}") print(f" 元素大小: {info['element_size']} bytes ({info['dtype']})") print(f" 形状: {info['shape']}") ![配图](https://i-blog.csdnimg.cn/img_convert/bdd943e5dbffa2bc176d36fbebd86536.png) print(f" 设备: {info['device']}") if 'gpu_allocated' in info: print(f" GPU 已分配: {info['gpu_allocated'] / 1e6:.2f} MB") print(f" GPU 已保留: {info['gpu_reserved'] / 1e6:.2f} MB") def demo_basic_memory(): """演示基本内存测量""" print("=" * 60) print("Tensor 内存测量对比") print("=" * 60) # 小 Tensor small = torch.randn(10, 10) print_tensor_memory(small, "small (10x10)") # 大 Tensor large = torch.randn(1000, 1000) print_tensor_memory(large, "large (1000x1000)") # 超大 Tensor huge = torch.randn(10000, 10000) print_tensor_memory(huge, "huge (10000x10000)") print() def demo_storage_sharing(): """演示 Storage 共享""" print("=" * 60) print("Storage 共享演示") print("=" * 60) original = torch.randn(1000, 1000) print_tensor_memory(original, "original") # view 共享 Storage view = original.view(500, 2000) print_tensor_memory(view, "view") print(f" 共享 Storage: {original.data_ptr() == view.data_ptr()}") # slice 共享 Storage sliced = original[:500, :500] print_tensor_memory(sliced, "sliced") print(f" 共享 Storage: {original.data_ptr() == sliced.data_ptr()}") # transpose 共享 Storage transposed = original.t() print_tensor_memory(transposed, "transposed") print(f" 共享 Storage: {original.data_ptr() == transposed.data_ptr()}") # clone 不共享 cloned = original.clone() print_tensor_memory(cloned, "cloned") print(f" 共享 Storage: {original.data_ptr() == cloned.data_ptr()}") print() def demo_dtype_memory(): """演示不同 dtype 的内存差异""" print("=" * 60) print("不同 dtype 的内存对比") print("=" * 60) shape = (1000, 1000) dtypes = [torch.float32, torch.float64, torch.float16, torch.int32, torch.int64, torch.int16, torch.int8, torch.bool] print(f" {'dtype':<15} {'element_size':>12} {'total_size':>15}") print(f" {'-'*15} {'-'*12} {'-'*15}") for dtype in dtypes: tensor = torch.zeros(shape, dtype=dtype) es = tensor.element_size() total = es * tensor.nelement() print(f" {str(dtype):<15} {es:>10} B {total:>12} B ({total/1e6:.2f} MB)") print() def demo_process_memory(): """使用 psutil 监控进程内存""" print("=" * 60) print("进程内存监控") print("=" * 60) process = psutil.Process(os.getpid()) mem_before = process.memory_info().rss / 1e6 print(f" 分配前 RSS: {mem_before:.2f} MB") # 分配大 Tensor tensors = [] for i in range(5): t = torch.randn(2000, 2000) # ~16MB each tensors.append(t) mem_now = process.memory_info().rss / 1e6 tensor_mem = sum(t.element_size() * t.nelement() for t in tensors) / 1e6 print(f" 分配 {i+1} 个后: RSS={mem_now:.2f} MB, Tensor数据={tensor_mem:.2f} MB") # 释放 del tensors mem_after = process.memory_info().rss / 1e6 print(f" 释放后 RSS: {mem_after:.2f} MB") print() def demo_gpu_memory(): """GPU 显存监控""" print("=" * 60) print("GPU 显存监控") print("=" * 60) if not torch.cuda.is_available(): print(" CUDA 不可用，跳过 GPU 演示") print() return torch.cuda.reset_peak_memory_stats() mem_before = torch.cuda.memory_allocated() / 1e6 print(f" 分配前 GPU 已分配: {mem_before:.2f} MB") # 分配 GPU Tensor tensors = [] for i in range(5): t = torch.randn(2000, 2000, device='cuda') # ~16MB each tensors.append(t) mem_now = torch.cuda.memory_allocated() / 1e6 print(f" 分配 {i+1} 个后: GPU已分配={mem_now:.2f} MB") peak = torch.cuda.max_memory_allocated() / 1e6 print(f" GPU 峰值显存: {peak:.2f} MB") # 释放 del tensors torch.cuda.empty_cache() mem_after = torch.cuda.memory_allocated() / 1e6 print(f" 释放后 GPU 已分配: {mem_after:.2f} MB") print() def demo_memory_estimation(): """模型内存估算工具""" print("=" * 60) print("模型内存估算") print("=" * 60) model = torch.nn.Sequential( torch.nn.Linear(1000, 2000), torch.nn.ReLU(), torch.nn.Linear(2000, 1000), torch.nn.ReLU(), torch.nn.Linear(1000, 100), ) # 参数内存 param_mem = 0 for name, param in model.named_parameters(): size = param.element_size() * param.nelement() param_mem += size print(f" {name}: {param.shape}, {size/1e6:.2f} MB") # 梯度内存（与参数相同） grad_mem = param_mem # 反向传播时 # 优化器状态（Adam: 2倍参数大小） optimizer_mem = param_mem * 2 # Adam 的一阶和二阶动量 # 激活值内存（取决于 batch_size） batch_size = 32 input_size = 1000 activation_mem = batch_size * (input_size + 2000 + 1000 + 100) * 4 # float32 total_train = param_mem + grad_mem + optimizer_mem + activation_mem print(f"\n 参数内存: {param_mem/1e6:.2f} MB") print(f" 梯度内存: {grad_mem/1e6:.2f} MB") print(f" 优化器内存 (Adam): {optimizer_mem/1e6:.2f} MB") print(f" 激活值内存 (batch={batch_size}): {activation_mem/1e6:.2f} MB") print(f" 训练总估算: {total_train/1e6:.2f} MB") print() if __name__ == '__main__': demo_basic_memory() demo_storage_sharing() demo_dtype_memory() demo_process_memory() demo_gpu_memory() demo_memory_estimation() print("=" * 60) print("关键总结:") print("1. sys.getsizeof 只测量 Python 对象，不含数据") print("2. 使用 element_size() * numel() 计算实际数据大小") print("3. view/slice/transpose 共享 Storage，不额外占内存") print("4. GPU 显存用 torch.cuda.memory_allocated() 监控") print("5. 训练内存 = 参数 + 梯度 + 优化器状态 + 激活值")

常见陷阱与注意事项

1. sys.getsizeof 不测量底层数据

# sys.getsizeof 只测量 Python 对象头，不包含 Storage 数据 tensor = torch.randn(1000000) print(sys.getsizeof(tensor)) # ~72 bytes，不是 4MB # 正确方式 print(tensor.element_size() * tensor.numel()) # 4000000 bytes

2. view 不复制数据

a = torch.randn(1000, 1000) b = a.view(500, 2000) # 共享 Storage，不额外占内存 c = a.t() # 共享 Storage d = a[:500] # 共享 Storage # 只有 clone() / contiguous() 可能复制数据

3. GPU 显存不释放

# 删除 Tensor 后 GPU 显存可能不立即释放 del gpu_tensor torch.cuda.empty_cache() # 手动触发垃圾回收

4. 梯度和优化器状态占用

# 训练时的实际内存远大于模型参数 # 参数: P bytes # 梯度: P bytes（反向传播时） # Adam 优化器: 2P bytes（一阶+二阶动量） # 总计: ~4P bytes（不含激活值）

5. 激活值内存

# 前向传播保存的中间激活值用于反向传播 # 激活值内存 = sum(batch_size * layer_output_size * element_size) # 使用 gradient checkpointing 可以减少激活值内存

6. 内存碎片

# 频繁分配/释放不同大小的 Tensor 会导致内存碎片 # CPU: 使用 psutil 观察 RSS vs 实际使用 # GPU: torch.cuda.memory_reserved() vs torch.cuda.memory_allocated()

7. pinned memory

# pin_memory=True 的 Tensor 使用锁页内存 # 不被操作系统换出，加速 CPU→GPU 传输 # 但会减少可用于其他程序的内存

总结

理解 PyTorch Tensor 内存占用的关键要点：

Python 对象 vs 底层数据：sys.getsizeof()只测量 Python 对象头（~72 bytes），不包含Storage中的实际数据
正确测量方法：tensor.element_size() * tensor.nelement()或tensor.nbytes（PyTorch 2.0+）
Storage 共享：view、slice、transpose创建的 Tensor 共享底层Storage，不额外占用数据内存
GPU 显存监控：使用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()
训练内存估算：参数 + 梯度 + 优化器状态（Adam 为 2 倍参数）+ 激活值
dtype 影响：float64是float32的 2 倍，float16是float32的一半
内存释放：del tensor+torch.cuda.empty_cache()释放 GPU 显存
进程内存：使用psutil.Process().memory_info().rss监控 CPU 进程内存

核心原则：不要用sys.getsizeof()测量 Tensor 的数据大小，使用element_size() * numel()或nbytes属性获取真实内存占用。

企业官网建设流程全解析