PyTorch推理模式与显存优化实战：从eval到inference_mode

推理部署时，模型能输出结果并不等于代码已经进入正确推理状态。深度学习项目若仍保留训练图、随机层或不必要的中间张量，会造成显存升高、延迟波动，甚至同一输入得到不同结果。本文梳理PyTorch推理模式的正确用法。

一、问题背景

常见误区是只调用model.eval()却忘记关闭梯度，或只使用torch.no_grad()却让Dropout保持训练状态。前者影响内存，后者影响结果稳定性。对于大模型训练后的验证，两项设置承担不同职责。

GPU服务器租用可以快速获得测试设备，但上线前仍要在目标批次、输入尺寸和精度下建立基线。AI算力平台上的空闲显存也不能直接等同于可承载并发，因为框架缓存、临时张量和输入峰值都会参与占用。

二、环境准备

准备CUDA版PyTorch、训练权重和测试样本：

nvidia-smi
python -c "import torch; print(torch.__version__)"

加载模型时先映射到CPU，再移动到GPU：

state = torch.load("model.pth", map_location="cpu")
model.load_state_dict(state)
model = model.cuda().eval()

临时压测可使用GPU算力平台。润云智算官网（https://www.smoothcloud.com.cn/）提供按需GPU资源和开发镜像，已确认环境包含Ubuntu、Python、CUDA、JupyterLab与SSH；实际版本应与模型依赖匹配。

三、编号实操步骤

1. 区分eval()与禁用梯度

eval()会切换Dropout、BatchNorm等模块行为，但不会关闭自动求导。完整基础写法如下：

model.eval()
with torch.no_grad():
output = model(inputs.cuda())

如果后处理把GPU张量长期保存在列表中，即使不记录梯度也会占用显存，应及时转到CPU或只保留必要结果。

2. 使用inference_mode

纯推理流程可尝试更严格的模式：

model.eval()
with torch.inference_mode():
output = model(inputs.cuda(non_blocking=True))

它除关闭梯度外，还减少部分张量版本跟踪开销。但其张量不适合随后重新进入需要自动求导的流程；若业务包含梯度解释或在线优化，应使用no_grad()并明确边界。

3. 记录峰值显存

torch.cuda.reset_peak_memory_stats()
with torch.inference_mode():
output = model(inputs.cuda())
torch.cuda.synchronize()

peak = torch.cuda.max_memory_allocated() / 1024**2
reserved = torch.cuda.memory_reserved() / 1024**2
print(f"peak={peak:.1f}MB reserved={reserved:.1f}MB")

allocated表示张量实际占用，reserved包含缓存分配器保留空间。不能看到reserved较高就认定显存泄漏。

4. 做预热与延迟测试

首次执行包含CUDA初始化和内核加载，不应计入稳定延迟：

import time

for _ in range(10):
with torch.inference_mode():
model(inputs.cuda())

torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(100):
with torch.inference_mode():
model(inputs.cuda())
torch.cuda.synchronize()
print((time.perf_counter() – start) / 100)

正式测试应报告平均值和高分位延迟，并固定输入、批次与精度。

5. 控制输出生命周期

需要长期保存结果时立即脱离GPU：

result = output.float().cpu().numpy()
del output

不要每次请求都调用torch.cuda.empty_cache()，它不能释放仍被引用的张量，还可能破坏缓存复用并造成延迟抖动。

6. 验证结果一致性

用同一批样本比较训练框架验证结果和服务结果，检查预处理、数据类型、类别映射及容差。模型微调结束后，也应先通过固定回归集，再进入推理部署压测。

四、常见问题与解决方案

1. 调用eval()后显存仍高

eval()不负责关闭梯度。检查是否使用no_grad()或inference_mode()，以及输出是否仍被容器引用。

2. 同一输入结果不同

确认模型已进入评估状态，同时检查随机预处理、采样参数和未固定的随机种子。

3. empty_cache()后显存没有下降

仍被Python变量引用的张量无法释放。先定位引用，再考虑清理缓存。

4. 批次增大后延迟突然上升

可能触发显存峰值、数据复制或算子路径变化。逐级压测批次，而不是直接使用最大值。

五、总结

稳定推理需要同时处理模块状态、梯度记录、张量生命周期和计时方法。eval()与inference_mode()并非替代关系，而是分别控制模型行为与执行开销。GPU算力平台能提供弹性测试资源，最终规格仍应由真实请求压测决定。

FAQ

Q1：inference_mode()一定比no_grad()快吗？

不一定。收益取决于模型结构和工作负载，应在目标环境中测试。

Q2：评估验证集也能使用inference_mode()吗？

如果不需要梯度或梯度解释，可以使用；需要反向计算时则不能使用。

Q3：如何判断显存是否泄漏？

观察多轮请求后memory_allocated是否持续增长，并检查保存输出、闭包和全局容器中的引用。

Q4：显存足够就能提高并发吗？

不一定。并发还受计算吞吐、CPU预处理、网络和响应时延目标影响。