ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

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

2026/8/30 16:06:12 拓冰建站 浏览量
【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(fTensor 元素数量: {big_tensor.numel()}) print(fTensor 实际数据大小: {big_tensor.element_size() * big_tensor.numel() / 1e9:.2f} GB) print(fsys.getsizeof(tensor): {sys.getsizeof(big_tensor)} bytes) # 输出可能只有几十字节为什么sys.getsizeof()报告的内存如此之小Tensor 的实际数据存储在哪里如何正确测量 Tensor 的内存占用错误复现场景一误用 sys.getsizeofimport torch import sys tensor torch.randn(1000, 1000) print(fsys.getsizeof: {sys.getsizeof(tensor)}) # 可能只有 72 bytes # 开发者误以为 Tensor 只占 72 字节实际上数据占 1000*1000*4 4MB场景二误用sizeoftensor 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(fStorage 大小: {tensor.storage().nbytes()} bytes) # 注意storage().nbytes() 返回的是 Storage 的总大小 # 可能大于 Tensor 实际使用的部分如有 padding方案三使用 nbytes 属性PyTorch 2.0tensor torch.randn(1000, 1000) print(fnbytes: {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, devicecuda) # 分配后 mem_after torch.cuda.memory_allocated() print(fGPU 显存增量: {(mem_after - mem_before) / 1e6:.2f} MB) print(fGPU 峰值显存: {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, nametensor): 打印 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, dtypedtype) 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 分配 {i1} 个后: 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, devicecuda) # ~16MB each tensors.append(t) mem_now torch.cuda.memory_allocated() / 1e6 print(f 分配 {i1} 个后: 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 bytes2. 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_memoryTrue 的 Tensor 使用锁页内存 # 不被操作系统换出加速 CPU→GPU 传输 # 但会减少可用于其他程序的内存总结理解 PyTorch Tensor 内存占用的关键要点Python 对象 vs 底层数据sys.getsizeof()只测量 Python 对象头~72 bytes不包含Storage中的实际数据正确测量方法tensor.element_size() * tensor.nelement()或tensor.nbytesPyTorch 2.0Storage 共享view、slice、transpose创建的 Tensor 共享底层Storage不额外占用数据内存GPU 显存监控使用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()训练内存估算参数 梯度 优化器状态Adam 为 2 倍参数 激活值dtype 影响float64是float32的 2 倍float16是float32的一半内存释放del tensortorch.cuda.empty_cache()释放 GPU 显存进程内存使用psutil.Process().memory_info().rss监控 CPU 进程内存核心原则不要用sys.getsizeof()测量 Tensor 的数据大小使用element_size() * numel()或nbytes属性获取真实内存占用。