PyTorch在Ubuntu上的内存管理优化策略
torch.cuda.empty_cache()可强制清理未使用的缓存，减少显存占用虚高（注意：此操作会触发同步，建议在调试阶段使用）。torch.cuda.memory_summary()查看显存分配详情（包括已用/预留显存、各张量占比）；torch.cuda.memory_allocated()和torch.cuda.memory_reserved()分别获取当前已分配和预留的显存大小，帮助快速定位内存瓶颈。del删除不再使用的变量。add_()、mul_()等原地操作减少新张量的创建（如x.add_(2)代替y = x.add(2)），降低显存占用。torch.no_grad()上下文管理器，避免计算图的生成（计算图会占用额外显存）。for i, (data, label) in enumerate(dataloader): ... if (i+1) % accumulation_steps == 0: scaler.step(optimizer); scaler.update(); optimizer.zero_grad()），在不增加显存的情况下提升训练稳定性。torch.cuda.amp模块（自动混合精度，AMP），结合float16（低精度）和float32（标准精度）计算，减少显存占用并加速训练（如scaler = torch.cuda.amp.GradScaler(); with torch.cuda.amp.autocast(): output = model(data); loss = criterion(output, label); scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()）。num_workers（数据加载的并行进程数）大于0（如num_workers=4），避免数据加载阻塞训练；设置pin_memory=True，将数据预加载到固定内存（Pinned Memory），加速数据从CPU到GPU的传输。yield）或迭代器逐批加载数据，避免一次性将所有数据读入内存（如def data_loader(file_path): with open(file_path, 'rb') as f: while True: data = f.read(64*1024); if not data: break; yield torch.from_numpy(np.frombuffer(data, dtype=np.float32))）。torch.autograd.profiler或NVIDIA Nsight Systems工具分析代码，定位未释放的张量或计算图（如循环中累积中间结果未用detach()断开引用）。例如，with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], profile_memory=True) as prof: ... print(prof.key_averages().table())。torch.utils.checkpoint在前向传播时丢弃不需要的中间激活，反向传播时重新计算，减少显存占用（如from torch.utils.checkpoint import checkpoint; def forward_with_checkpoint(segments): return checkpoint(segments)）。torch.backends.cudnn.benchmark=True，让CUDA自动选择最优的卷积算法（提升计算效率）；设置torch.backends.cudnn.deterministic=True，控制模型的随机性（便于调试）。os._exit(0)强制终止进程），彻底释放CUDA资源。