训练时CPU内存逐轮增长且不释放问题排查与解决请求

1. torch.load缓存未清理：torch.load依赖pickle加载数据，pickle会缓存已加载的对象，多次迭代加载时内存会持续累积；若未指定map_location，张量可能被加载到非预期设备，间接增加CPU内存占用。
2. collate_fn张量残留：collate_fn中提前将padding张量绑定到GPU设备时，主进程删除返回张量后，部分中间引用无法被GC及时回收；若在CPU设备操作，PyTorch的内存分配器会缓存已分配内存块，不会立即还给操作系统，导致内存占用看起来持续增长。
3. 自定义DataLoader的引用循环：自定义的VideoEmbeddingDataLoader同时持有self.dataset和父类DataLoader的dataset引用，形成循环引用，阻碍GC正常回收对象。

1. 优化torch.load的加载逻辑 #

修改__getitem__方法，指定加载位置并清理pickle缓存：

def __getitem__(self, index):
    tensor_path = self.pth_file[index]
    # 强制加载到CPU，避免设备不匹配导致的内存冗余
    with torch.no_grad():
        tensor = torch.load(tensor_path, map_location='cpu')
    # 清理pickle的加载缓存
    import pickle
    pickle.clear_cache()
    label = self.labels[index]
    return tensor, label

2. 调整collate_fn的张量处理 #

不在collate_fn中提前移动张量到GPU，改为训练时统一转移，减少中间变量残留：

def collate_fn(self, batch):
    max_frames = max([len(sample[0]) for sample in batch])
    padded_tensor = []
    padded_labels = []
    for tensor, label in batch:
        num_padding = max_frames - len(tensor)
        if num_padding > 0:
            # 先在CPU创建padding张量，后续统一移到目标设备
            padding_tensor = torch.zeros(num_padding, tensor.shape[1], dtype=tensor.dtype)
            tensor = torch.cat((tensor, padding_tensor), dim=0)
        else:
            tensor = tensor[:max_frames]
        padded_tensor.append(tensor)
        padded_labels.append(label)
    return torch.stack(padded_tensor, dim=0), padded_labels

3. 移除自定义DataLoader类 #

直接使用原生PyTorch DataLoader，避免不必要的引用循环：

# 删除自定义的VideoEmbeddingDataLoader类，替换为以下代码
batch_size = 2
# 训练集
train_dataset = VideoEmbeddingDataset(train_embedding_files, train_labels)
train_dataloader = DataLoader(
    train_dataset, 
    batch_size=batch_size, 
    shuffle=True, 
    num_workers=0, 
    collate_fn=train_dataset.collate_fn, 
    pin_memory=False
)
# 验证集
val_dataset = VideoEmbeddingDataset(val_embedding_files, val_labels)
val_dataloader = DataLoader(
    val_dataset, 
    batch_size=batch_size, 
    shuffle=False, 
    num_workers=0, 
    collate_fn=val_dataset.collate_fn, 
    pin_memory=False
)

4. 强化内存回收流程 #

在迭代循环中，除了删除张量和GC回收，还需清理PyTorch内存分配器缓存：

import gc
import torch
from tqdm import tqdm

num_epochs = 5
for epoch in range(num_epochs):
    print(f'Epoch {epoch+1}/{num_epochs}')
    for embeddings_tensor, batch_labels in tqdm(train_dataloader):
        # 强制释放当前批次张量
        del embeddings_tensor
        del batch_labels
        # 清理PyTorch CPU/GPU内存缓存
        torch.memory.empty_cache()
        torch.cuda.empty_cache()
        # 触发Python垃圾回收
        gc.collect()

5. 调整num_workers参数（可选） #

若num_workers=0导致主进程内存累积，可尝试设置num_workers为CPU核心数，让数据加载在子进程执行，子进程结束后自动释放内存：

train_dataloader = DataLoader(
    train_dataset, 
    batch_size=batch_size, 
    shuffle=True, 
    num_workers=4,  # 根据CPU核心数调整
    collate_fn=train_dataset.collate_fn, 
    pin_memory=True
)

内容的提问来源于stack exchange，提问作者Rabeya Akter