先聊聊你遇到的那个奇怪现象：加了print(torch.cuda.is_available())就临时好了，其实是因为这个调用会主动触发PyTorch初始化CUDA上下文——有时候PyTorch延迟初始化CUDA时，内存分配逻辑可能出现异常，提前触发初始化相当于把这个潜在问题给“抹平”了，但这只是临时 workaround，不是长久之计。下面给你几个能彻底解决问题的方案：

1. 精准定位内存占用源头 #

训练还没启动就爆内存，大概率是初始化阶段偷偷创建了不必要的CUDA张量。你可以在代码关键节点插入内存监控代码，找出内存泄漏的位置：

import torch

# 在导入后、模型初始化前、数据加载前分别打印
print(f"已分配内存: {torch.cuda.memory_allocated()/1024**2:.2f} MiB")
print(f"已预留内存: {torch.cuda.memory_reserved()/1024**2:.2f} MiB")

重点检查：有没有在数据预处理时把整个数据集移到GPU？模型初始化时有没有超大参数层？第三方库有没有偷偷占用GPU内存？

2. 优化PyTorch内存分配策略 #

PyTorch默认的内存分配器有时候会产生碎片，导致明明有剩余内存却无法分配。可以通过环境变量调整：

import os
# 开启可扩展内存段，减少碎片
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

另外，你可以限制当前进程的GPU内存使用率，避免和残留进程抢占资源：

# 限制当前进程使用90%的GPU 0内存
torch.cuda.set_per_process_memory_fraction(0.9, device=0)

3. 彻底清理GPU残留进程 #

重启Jupyter内核有时候并不能完全杀死GPU上的残留进程——尤其是之前的代码崩溃后，可能有僵尸进程占着内存。你可以用命令行工具手动清理：

# 查看GPU占用情况，找到占用内存的进程PID
nvidia-smi
# 杀死对应进程，替换<pid>为实际进程号
kill -9 <pid>

如果用的是Jupyter，建议关闭所有无关的笔记本，避免多个进程共享GPU内存。

4. 轻量化模型与数据处理 #

如果你的模型或batch size太大，初始化阶段就会占满内存：

• 减小batch size：这是最直接的方法，先从极小的batch size（比如2、4）开始测试，逐步找到最大可行值。
• 启用半精度训练：用torch.cuda.amp.autocast()包裹训练循环，能在几乎不损失精度的情况下把内存占用降低一半左右。
• 检查模型结构：有没有可以简化的层？比如把大的全连接层换成注意力机制，或者用更小的特征图尺寸。

5. 主动初始化CUDA上下文 #

把torch.cuda.empty_cache()放在导入后还不够，你需要先确保CUDA上下文已经正确初始化，再清理缓存：

import torch
# 主动初始化CUDA上下文
torch.cuda.init()
# 再清空缓存
torch.cuda.empty_cache()

这比单纯调用empty_cache()更可靠，因为它确保缓存清理是在CUDA环境完全就绪的状态下执行的。

内容的提问来源于stack exchange，提问作者Debvrat Varshney