深度学习推理显存泄露排查与处理实操指南

樽酒ﻬق

已于 2025-04-04 10:46:32 修改

阅读量1.1k

点赞数 23

分类专栏：机器学习文章标签：深度学习人工智能

于 2025-04-04 10:42:15 首次发布

本文链接：https://blog.csdn.net/u010702254/article/details/146995629

版权

机器学习专栏收录该内容

8 篇文章

订阅专栏

深度学习推理显存泄露排查与处理实操指南

1. 显存泄露原理与成因

1.1 显存泄露基本概念

GPU显存泄露是指在模型训练或推理过程中，由于代码实现不当或框架缺陷，导致GPU显存被占用后无法正常释放，从而造成显存占用量持续累积上升的现象。与CPU内存泄露类似，但显存泄露专指GPU上的显存资源被不断消耗，最终导致"Out of Memory"(OOM)错误，使模型无法继续运行。

在深度学习推理系统中，显存泄露尤为致命，因为推理服务通常需要长时间稳定运行，即使极小的泄露也会随着时间积累而导致系统崩溃。

1.2 常见泄露原因分析

根据泄露来源，我们可以将显存泄露原因分为以下几类：

模型规模问题：模型过大、参数量超出GPU显存容量。
批量大小设置不当：推理时batch size设置过大，每次处理的数据量超出显存承载能力。
数据加载和处理问题：数据加载方式不当导致张量在GPU上累积，如将过多的预处理数据转移到GPU。
显存管理问题：最常见的泄露原因，包括：
- 未正确释放不再使用的张量（Tensor）
- 张量引用被意外保存在长期存在的数据结构中
- 显存碎片化问题
- 计算图构建不当导致的累积
框架特定问题：不同深度学习框架的内部实现可能导致特定类型的显存泄露。

1.3 不同框架中的典型泄露模式

PyTorch中的典型泄露

PyTorch中的显存泄露主要源于以下几点：

Tensor的requires_grad与计算图：当Tensor的requires_grad为True时，每次运算都会自动构建计算图，保存AutogradMeta信息，如果不断进行运算而不释放，就会导致显存泄露fatescript.github.io。
循环引用：PyTorch中模型组件之间的循环引用可能导致对象无法被垃圾回收，其占用的显存也无法释放。
DataLoader相关问题：数据加载过程中，如果相关张量被转移到GPU而没有适当释放，会导致显存累积。

TensorFlow中的典型泄露

TensorFlow中常见的显存泄露问题包括：

运行时环境问题：TensorFlow Sessions不正确关闭或重用，导致图计算时分配的显存无法及时释放。
计算图构建问题：在循环中反复创建新的计算图节点而不释放，导致显存不断累积。
特征顺序变化问题：TensorFlow Serving中，由于PredictRequest中inputs特征顺序变化，导致DirectSession不断创建新的存储结构，形成显存泄露。
缓存累积问题：TensorFlow自动缓存某些计算结果以提高性能，但在长时间运行的推理服务中可能导致缓存过度增长。

2. 泄露检测与分析工具

2.1 通用监控工具

基础监控工具是发现显存泄露的第一道防线：

nvidia-smi：最基本的NVIDIA GPU监控工具，可以通过命令watch -n 1 nvidia-smi实时观察显存使用情况。
NVTOP：类似于top命令的交互式GPU监控工具，可以实时查看GPU使用率和显存占用。
Python监控函数：自定义函数定期检查GPU显存使用情况：

def log_gpu_memory_usage():
    """记录GPU显存使用情况"""
    import torch
    
    gpu_allocated = torch.cuda.memory_allocated() / (1024 * 1024)  # MB
    gpu_reserved = torch.cuda.memory_reserved() / (1024 * 1024)  # MB
    gpu_max_allocated = torch.cuda.max_memory_allocated() / (1024 * 1024)  # MB
    
    print(f"GPU Allocated: {gpu_allocated:.2f} MB, Reserved: {gpu_reserved:.2f} MB, Max Allocated: {gpu_max_allocated:.2f} MB")

2.2 PyTorch专用工具

PyTorch提供了多种高级工具用于分析显存使用：

基础显存查询函数：
- torch.cuda.memory_allocated()：显示当前GPU上张量实际使用的显存
- torch.cuda.memory_reserved()：显示PyTorch为操作保留的显存
- torch.cuda.max_memory_allocated()：记录峰值显存使用Medium

显存快照工具（Memory Snapshot）：

# 开始记录显存事件
torch.cuda.memory._record_memory_history(max_entries=100000)

# 模型训练或推理代码
# ...

# 保存显存快照
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# 停止记录
torch.cuda.memory._record_memory_history(enabled=None)

快照文件可以使用在线工具（https://pytorch.ac.cn/memory_viz）进行可视化分析PyTorch中文网。

显存分析器（Memory Profiler）：

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    # 模型推理代码
    # ...
    prof.step()

# 导出显存时间线
prof.export_memory_timeline("memory_timeline.html", device="cuda:0")

2.3 TensorFlow工具

TensorFlow生态提供的显存分析工具：

tf.debugging.set_log_device_placement：跟踪操作执行的设备和显存分配情况。
TensorBoard的Profiler：可以分析模型运行时的显存消耗并可视化。

tf.config.experimental.get_memory_info：获取当前可用和总显存信息。

import tensorflow as tf
memory_info = tf.config.experimental.get_memory_info('GPU:0')
print(f"Available: {memory_info['current'] / (1024**2)} MB")
print(f"Total: {memory_info['peak'] / (1024**2)} MB")

2.4 NVIDIA专业工具

NVIDIA提供的专业调试与监控工具：

NVIDIA Compute Sanitizer：
- memcheck：检测显存访问错误和泄露
- racecheck：检测共享显存数据访问危险
- initcheck：检测未初始化的显存访问
- synccheck：线程同步危险检测NVIDIA Developer
使用方法示例：
```
compute-sanitizer --tool memcheck --leak-check=full ./your_executable
```
NVIDIA Nsight Systems：系统级性能分析工具，可视化程序执行时间线和显存使用情况。
NVIDIA Nsight Compute：针对CUDA内核性能和显存分析的工具。

通过上述工具，可以全方位地监控、检测和分析GPU显存使用情况，及时发现潜在的泄露问题。

3. 预防与处理策略

3.1 代码层面优化

合理使用requires_grad：在进行推理时，使用torch.no_grad()上下文管理器或设置requires_grad=False，避免不必要的计算图构建导致显存泄露fatescript.github.io。
```
# 推理时关闭梯度计算
with torch.no_grad():
    outputs = model(inputs)

# 或手动设置requires_grad=False
inputs = inputs.requires_grad_(False)
```

手动释放不再使用的GPU张量：

# 在不需要时主动删除变量
del tensor_on_gpu
# 清理CUDA缓存
torch.cuda.empty_cache()

注意避免循环引用：检查模型或数据结构中是否存在相互引用的对象，这可能阻止垃圾回收器回收GPU上的张量。
正确处理自定义数据：确保在数据加载和预处理过程中不会意外地将过多数据持久化到GPU显存中CSDN。
批量处理数据，设置合理的batch size：在推理时，考虑将输入数据分成较小的批次处理，避免单批次数据过大占用过多显存。

3.2 框架级优化策略

PyTorch优化：

定期使用torch.cuda.empty_cache()清理缓存显存
使用torch.cuda.synchronize()确保GPU操作完成并正确释放临时显存
考虑设置环境变量PYTORCH_NO_CUDA_MEMORY_CACHING=1禁用显存缓存（但可能影响性能）

# 例如，在批处理循环中：
for i, batch in enumerate(dataloader):
    # 处理批次
    outputs = model(batch.to('cuda'))
    # 处理输出...
    
    # 每N个批次清理一次显存
    if i % N == 0:
        # 确保所有CUDA操作完成
        torch.cuda.synchronize()
        # 释放缓存显存
        torch.cuda.empty_cache()

TensorFlow优化：
- 使用tf.config.experimental.set_memory_growth(gpu, True)允许TensorFlow按需分配显存
- 在TensorFlow 2.x中使用tf.keras.backend.clear_session()释放模型占用的显存
```
# 控制显存增长
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
```
ONNX Runtime优化：
- 利用内存优化选项减少显存占用
- 开启显存复用功能（memory pattern和memory reuse）

3.3 模型优化方法

模型压缩与量化：
- 使用模型剪枝减少参数数量，直接降低显存需求
- 采用权重量化（如FP16或INT8量化）减少显存占用
- 知识蒸馏生成更小的等效模型
中间结果优化：
- 减少或重用中间激活值存储
- 优化模型结构，减少中间特征图尺寸和数量
模型结构优化：
- 使用更高效的网络架构
- 减少网络深度或宽度
- 使用TensorRT等工具进行算子融合，减少显存占用

3.4 部署架构优化

美团视觉团队的实践表明，合理的推理服务架构设计能显著提高GPU利用效率并减轻显存压力[参考地址]：

功能分离：将模型的CPU运算部分（如图像预处理、后处理）与GPU运算部分（主干网络）分离，分别部署为独立的微服务，减少GPU上的非核心计算，释放显存。
异步处理：使用异步处理机制，避免在GPU上积累过多等待处理的数据。
资源分配：根据负载情况动态调整服务的实例数量，确保单个GPU不会过载导致显存问题。
批处理策略：实现智能批处理策略，在保证低延迟的同时最大化吞吐量，减少因频繁启动推理而导致的显存碎片。

通过这种分离式架构，可以显著减少单个请求处理过程中的显存占用，同时提高整体推理效率。

4. 实战案例分析

4.1 PyTorch中的实例分析

案例一：DataLoader引起的显存泄露

在一个物体检测模型训练中，开发者发现即使注释掉模型训练代码，仅保留数据加载部分，显存仍然持续上涨[参考地址]。

问题代码：

for e in range(epoch):
    for i, (imgs, targets) in enumerate(train_dataloader):
        imgs = [img.to(device) for img in imgs]
        targets = [_to_device(target, device) for target in targets]
        # 模型训练代码已注释

分析过程：

通过watch -n 1 nvidia-smi监控显存使用，发现持续上涨。
逐步注释数据的各个部分（imgs、boxes、labels、masks），发现当masks被注释时，显存不再泄露。
进一步分析发现，问题出在targets字典中的masks数据上，其中的tensor被长期保留在GPU显存中。

解决方案：

for e in range(epoch):
    for i, (imgs, targets) in enumerate(train_dataloader):
        imgs = [img.to(device) for img in imgs]
        targets = [_to_device(target, device) for target in targets]
        # 模型训练代码
        
        # 批次处理完后，删除targets中的masks并清理缓存
        for target in targets:
            del target['masks']
        torch.cuda.empty_cache()

案例二：requires_grad导致的计算图积累

在一个神经网络推理服务中，开发者发现即使是简单的前向传播操作，显存也会随着时间持续增长[参考地址]。

问题代码：

def inference(model, data):
    output = model(data)
    return output

# 在循环中重复调用
for _ in range(iterations):
    result = inference(model, input_data.to('cuda'))
    # 处理结果...

分析过程：

使用PyTorch的显存监控工具记录显存使用，发现持续上涨。
创建最小复现代码，确认当tensor的requires_grad为True时，显存不断增加。
分析原因：每次运算都会自动记录梯度计算所需的信息到计算图中，这些计算图没有被及时释放。

解决方案：

def inference(model, data):
    with torch.no_grad():  # 使用no_grad上下文
        output = model(data)
    return output

# 或者
def inference(model, data):
    output = model(data)
    return output.detach()  # 分离计算图

4.2 TensorFlow Serving案例

爱奇艺深度学习平台团队发现了两个严重的TensorFlow Serving显存泄露问题CSDN：

问题一：输入特征顺序导致的显存泄露

现象：TF Serving容器内存持续增长且没有停止迹象，GPU显存也不断增加。

分析过程：

使用gperftools进行内存profiling，间接发现显存异常增长的原因。
发现DirectSession::GetOrCreateExecutors函数创建了大量String对象，这些对象对应GPU上的显存分配。
深入分析发现原因是PredictRequest中inputs特征顺序不断变化，导致内部哈希表不断创建新条目，与之对应的GPU计算资源也不断增加。

解决方案：

对PredictRequest中的inputs进行排序，确保相同顺序。
修改TensorFlow源码中GetOrCreateExecutors函数逻辑。
提交两个PR分别修复TensorFlow和TF Serving中的显存泄露问题。

问题二：高并发下GRPC线程数量爆炸导致显存耗尽

现象：服务突增高并发请求后，TF Serving容器中grpcpp_sync_ser线程数量不断增加，这些线程会分配大量GPU计算资源，最终导致显存耗尽。

分析过程：

分析发现GRPC Server不断创建新的WorkerThread处理队列消息，但这些线程没有及时销毁。
每个线程都会为GPU任务分配资源，导致显存不断增加。

解决方案：
设置Resource Quota限制GRPC最大线程数量：

ServerBuilder builder;
grpc::ResourceQuota quota;
quota.SetMaxThreads(256);  // 设置最大线程数为256
builder.SetResourceQuota(quota);

4.3 ONNX Runtime优化

ONNX Runtime作为一个通用的高性能推理引擎，提供了多种优化选项来加速模型推理并减少显存占用[参考地址]：

PyTorch模型转ONNX并优化显存使用：

import torch
import torch.onnx

# 准备模型和样例输入
model = YourModel()
model.eval()
dummy_input = torch.randn(1, 3, 224, 224, device='cuda')

# 导出ONNX模型
torch.onnx.export(
    model, 
    dummy_input, 
    "model.onnx", 
    opset_version=12, 
    input_names=["input"], 
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

使用ONNX Runtime优化显存：

import onnxruntime as ort

# 创建优化的会话
options = ort.SessionOptions()
# 启用显存优化
options.enable_mem_pattern = True
options.enable_mem_reuse = True
# 限制GPU显存增长
options.gpu_mem_limit = 1024 * 1024 * 1024  # 限制使用1GB显存

# 创建推理会话
session = ort.InferenceSession("model.onnx", options, providers=["CUDAExecutionProvider"])

# 执行推理
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: input_data})

优化措施包括：

启用显存模式优化，根据模型的显存访问模式分配和重用显存。
设置最大显存分配限制，防止显存使用超出预期。
使用动态批处理技术处理多个样本，减少显存浪费。

5. 最佳实践与经验总结

5.1 开发阶段预防措施

定期检查显存使用：
- 在开发过程中定期监控GPU显存使用情况
- 使用显存分析工具检测潜在的显存泄露
- 为关键函数编写显存测试用例
遵循显存管理最佳实践：
- 推理时使用torch.no_grad()或等效方法
- 避免在循环中重复创建GPU上的张量
- 正确处理批次间的中间变量，及时将不需要的数据从GPU移出
代码审查与测试：
- 对可能影响显存使用的代码进行专门的代码审查
- 编写长时间运行的测试，确保显存使用稳定

5.2 部署阶段应对策略

合理的资源配置：
- 根据模型大小和批处理需求选择合适的GPU
- 在容器或服务中设置适当的显存限制
- 使用虚拟化技术（如MPS或MIG）合理分配GPU显存
服务架构设计：
- 采用分离式架构设计，减少单个服务的显存占用
- 实现优雅的显存回收和服务重启机制
- 考虑使用模型分片或并行技术分散显存压力
容错与恢复机制：
- 实现显存使用监控和自动告警
- 设计服务自动重启或扩缩容机制，防止显存泄露导致服务中断