Python内存及时释放方案求助：PyTorch张量缓存内存未正常释放

问题核心 #

执行move_tensor后未释放预期的3GB内存，执行loading_tensor后内存占用升至8GB，完全不符合先释放再加载的设计预期。

原因分析 #

1. 切片操作的视图特性：PyTorch中张量切片[NUM:]默认返回原张量的视图而非独立副本，视图会保留对原张量数据的引用。这意味着move_tensor中赋值后，原大张量并未被垃圾回收（GC），内存仍被占用。
2. 内存回收时机延迟：Python的GC不会立即回收不再使用的对象，memory_profiler的函数级监控可能在GC触发前就统计了内存，导致看起来没释放。
3. cat操作的额外内存开销：torch.cat会创建新张量，此时原缓存张量（包括视图和未被回收的原大张量）仍在内存中，加上新加载的3GB数据，总占用就会达到8GB。

解决方法 #

1. 强制切断原张量引用并触发回收 #

修改move_tensor，将切片结果转为独立副本，同时手动触发GC和CUDA缓存清理（如果用GPU）：

def move_tensor():
    # 创建切片的独立副本，切断对原张量的引用
    self.cacheData[0] = self.cacheData[0][NUM:].contiguous()
    self.cacheData[1] = self.cacheData[1][NUM:].contiguous()
    
    # 手动触发垃圾回收
    import gc
    gc.collect()
    
    # GPU场景下清理CUDA缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

2. 排查引用泄漏 #

检查项目中是否有其他变量（比如全局变量、其他类实例）持有原self.cacheData张量的引用，这些引用会阻止GC回收内存。

3. 预分配张量优化内存使用 #

避免频繁用torch.cat创建新张量，初始化时预分配固定大小的张量，通过索引赋值更新：

# 初始化阶段预分配足够大的张量（根据业务需求设置MAX_SIZE）
MAX_SIZE = 1000000  # 示例值
self.cacheData[0] = torch.zeros(MAX_SIZE, dtype=torch.float32)
self.cacheData[1] = torch.zeros(MAX_SIZE, dtype=torch.float32)
self.current_len = 0  # 记录当前有效数据长度

def move_tensor():
    self.current_len -= NUM
    # 直接截断有效长度，无需修改张量本身（后续赋值覆盖旧数据）
    self.current_len = max(self.current_len, 0)

def loading_tensor():
    # 确保有足够空间容纳新数据
    assert self.current_len + len(tmp_data) <= MAX_SIZE
    # 通过索引赋值，避免创建新张量
    self.cacheData[0][self.current_len:self.current_len+len(tmp_data)] = tmp_data
    self.cacheData[1][self.current_len:self.current_len+len(tmp_data)] = tmp_data
    self.current_len += len(tmp_data)

4. 精准监控内存 #

使用PyTorch自带的内存工具替代memory_profiler，实时查看内存变化：

# CPU内存监控（需安装psutil）
import psutil
def get_cpu_memory():
    return psutil.Process().memory_info().rss / 1024**3  # 转换为GB

# GPU内存监控
def get_gpu_memory():
    if torch.cuda.is_available():
        return torch.cuda.memory_allocated() / 1024**3  # 转换为GB
    return 0

# 在关键步骤打印内存
def run():
    print(f"Before move: {get_gpu_memory():.2f}GB")
    move_tensor()
    print(f"After move: {get_gpu_memory():.2f}GB")
    loading_tensor()
    print(f"After load: {get_gpu_memory():.2f}GB")

内容的提问来源于stack exchange，提问作者YA xiang