23.67 GiB GPU训练PatchCore时遇OutOfMemoryError问题求助

问题背景 #

在23.67 GiB显存的GPU上训练输入尺寸128×512的PatchCore模型时遭遇CUDA OOM错误，环境为CUDA 12.4、PyTorch 2.5.1，已尝试调小batch size至1、处理内存碎片、排查进程占用但未解决。

PatchCore专属优化手段 #

• 更换轻量骨干网络：将当前的wide_resnet50_2替换为参数更少的模型（如resnet18、mobilenet_v3_large），宽残差网络的特征图尺寸和参数量远大于轻量模型，能直接降低特征提取阶段的内存占用。
• 减少特征提取层数量：当前配置使用layer2+layer3，可改为仅保留layer2，或替换为更浅的组合（如layer1+layer2），减少生成的特征图数量与尺寸，降低内存消耗。
• 降低Coreset采样比例：调小coreset_sampling_ratio（如从0.1改为0.05或0.02），该参数控制训练时保留的特征样本比例，比例越低，内存中存储的特征库体积越小。
• 减少邻居计算数量：降低num_neighbors（如从9改为5或3），PatchCore计算异常分数时需匹配特征邻居，减少邻居数量可降低相似度计算阶段的内存占用。

PyTorch与训练流程优化 #

• 启用混合精度训练：使用torch.cuda.amp.autocast()和GradScaler，将张量以半精度（FP16）存储，可大幅降低显存占用，Anomalib框架支持直接开启该模式。
• 关闭梯度计算：PatchCore训练阶段仅需提取正常样本特征构建特征库，无需反向传播更新骨干网络，可将骨干网络设为eval()模式，并用torch.no_grad()包裹特征提取代码，避免存储梯度占用内存。
• 调整数据加载配置：将num_workers从8降至4，减少CPU侧数据预处理对GPU内存的间接占用；若仍有压力，可关闭非必要的数据增强（如RandomAdjustSharpness）。
• 定期清理显存：在每个训练epoch结束后调用torch.cuda.empty_cache()，配合PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True进一步缓解内存碎片问题。

内容的提问来源于stack exchange，提问作者林芷翎