生产环境中TensorRT推理显存泄漏如何定位？从引擎序列化到显式资源释放的完整分析

一、什么是TensorRT推理的显存泄漏？

在生产环境用TensorRT做深度学习推理（比如图像分类、目标检测）时，常会遇到服务跑了几天，显卡显存从几百MB涨到几个GB，最后报显存不足（OOM），重启后才回落——这就是典型的TensorRT显存泄漏。简单说，就是程序用了显存但没还给系统，一直占着导致后续请求没法分配资源。

1.1 生产中常见的泄漏表现

最明显的是显存持续上涨，增长速度和请求量正相关，重启才会下降。比如每秒处理100张图的服务，正常显存稳定在1GB，跑一周后涨到8GB，且不会自动回落，这种情况不是模型太大，而是资源没释放。

二、从引擎序列化到资源释放的全流程踩坑点

TensorRT推理流程是：构建网络→创建引擎→序列化引擎（存文件，下次直接加载）→加载引擎→创建上下文→执行推理→释放资源，每个环节都可能留下未回收的显存，下面用代码示例说明。

2.1 引擎序列化时的资源残留

序列化引擎是为了加速启动，但很多人忽略构建网络的中间对象（比如builder、network、config）。这些是TensorRT的C++核心对象，Python垃圾回收不会自动释放，必须手动删除。

技术栈：TensorRT 8.6 + Python 3.10 + PyCUDA

错误示例：

import tensorrt as trt

# 错误：序列化后中间对象没释放，占着显存
def serialize_engine(model_path):
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    config = builder.create_builder_config()
    
    # （省略）填充network，比如解析ONNX模型
    serialized_engine = builder.build_serialized_network(network, config)
    
    with open("model.trt", "wb") as f:
        f.write(serialized_engine)
    
    # 忘记删除builder、network、config，它们的显存没释放
serialize_engine("model.onnx")

正确示例：

import tensorrt as trt

# 正确：序列化后手动删除所有中间TensorRT对象
def serialize_engine(model_path):
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    config = builder.create_builder_config()
    
    # （省略）填充network
    serialized_engine = builder.build_serialized_network(network, config)
    
    with open("model.trt", "wb") as f:
        f.write(serialized_engine)
    
    # 手动删除触发C++层析构，释放显存
    del config, network, builder, serialized_engine
serialize_engine("model.onnx")

关键：所有TensorRT核心对象必须手动del，因为GC管不到C++对象的内存。

2.2 推理时的上下文和显存泄漏

引擎加载后，要创建IExecutionContext执行推理，还要给输入输出分配GPU显存。很多服务每次请求都新建资源，用完不释放，导致显存持续上涨。

技术栈同上

错误示例：

import tensorrt as trt
import pycuda.driver as cuda
import numpy as np

class BadService:
    def __init__(self):
        # 加载引擎
        with open("model.trt", "rb") as f:
            engine_data = f.read()
        self.engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(engine_data)
    
    def infer(self, img):
        # 错误1：每次请求新建上下文，用完不删
        context = self.engine.create_execution_context()
        # 错误2：每次分配显存，用完不释放
        input_gpu = cuda.mem_alloc(img.nbytes)
        output_gpu = cuda.mem_alloc(1000*4)
        
        cuda.memcpy_htod(input_gpu, img)
        context.execute_v2(bindings=[int(input_gpu), int(output_gpu)])
        output = np.empty(1000, dtype=np.float32)
        cuda.memcpy_dtoh(output, output_gpu)
        
        # 完全没释放资源，每次请求都占显存
        return output

# 跑1000次请求，显存爆炸
bad_service = BadService()
for _ in range(1000):
    img = np.random.randn(3,224,224).astype(np.float32)
    bad_service.infer(img)

正确示例：

import tensorrt as trt
import pycuda.driver as cuda
import numpy as np

class GoodService:
    def __init__(self):
        # 引擎只加载1次，复用
        with open("model.trt", "rb") as f:
            engine_data = f.read()
        self.engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(engine_data)
        # 上下文只创建1次，复用
        self.context = self.engine.create_execution_context()
        # 预分配固定大小的显存，不用每次分配
        self.input_shape = (3,224,224)
        self.output_size = 1000
        self.input_gpu = cuda.mem_alloc(np.prod(self.input_shape)*4)
        self.output_gpu = cuda.mem_alloc(self.output_size*4)
    
    def infer(self, img):
        cuda.memcpy_htod(self.input_gpu, img)
        self.context.execute_v2(bindings=[int(self.input_gpu), int(self.output_gpu)])
        output = np.empty(self.output_size, dtype=np.float32)
        cuda.memcpy_dtoh(output, self.output_gpu)
        return output
    
    # 服务停止时强制释放所有资源
    def release(self):
        cuda.free(self.input_gpu)
        cuda.free(self.output_gpu)
        del self.context, self.engine

# 用try-finally确保资源一定释放
good_service = GoodService()
try:
    for _ in range(1000):
        img = np.random.randn(3,224,224).astype(np.float32)
        res = good_service.infer(img)
finally:
    good_service.release()

核心：复用上下文和显存，服务退出时必须释放资源，用try-finally确保逻辑执行。

2.3 临时显存的泄漏

TensorRT构建引擎或推理时会分配临时显存（比如workspace），如果是手动申请的临时buffer，必须手动cuda.free，否则也会泄漏。大部分情况下，删除核心对象会自动释放临时显存，但自定义层的临时内存要注意。

三、泄漏定位的实用方法

找到泄漏点是解决问题的关键，下面几个方法可快速定位：

3.1 监控显存实时变化

用PyCUDA的mem_get_info()查看已用显存，每次推理前后调用，看是否上涨：

import pycuda.driver as cuda

def print_gpu_usage():
    free, total = cuda.mem_get_info()
    used = (total - free)/(1024**2) # 转MB
    print(f"已用：{used:.2f} MB，空闲：{free/(1024**2):.2f} MB")

# 测试时的打印
print("启动前：")
print_gpu_usage()
service = GoodService()
print("预分配后：")
print_gpu_usage()
for i in range(5):
    service.infer(np.random.randn(3,224,224).astype(np.float32))
    print(f"第{i+1}次推理后：")
    print_gpu_usage()
service.release()
print("释放后：")
print_gpu_usage()

如果每次推理后显存涨，说明infer里的资源没释放；服务释放后显存回落到初始值，说明资源管理正确。

3.2 拆分流程缩小范围

把代码拆成小部分测试：先测序列化引擎，再测加载，最后测推理。比如序列化时显存没释放，就是2.1的问题；推理时涨就是2.2的问题，快速缩小范围。

3.3 查看TensorRT详细日志

把日志级别设为VERBOSE，会打印资源分配和释放记录，找只出现“Allocated”没“Released”的buffer，就是泄漏点：

# 把logger级别设为VERBOSE，替换之前的WARNING
logger = trt.Logger(trt.Logger.VERBOSE)

四、技术优缺点和注意事项

4.1 资源管理方式的优缺点

• 手动删除对象+try-finally：优点是可控，确保资源释放；缺点是需要写额外代码，生产环境必须。
• 复用资源：优点是减少内存开销，提升性能；缺点是要固定输入形状，动态形状需调整。

4.2 生产环境核心注意事项

1. 不要每次请求新建引擎/上下文：引擎加载慢，复用可减少泄漏和性能损耗。
2. 服务退出强制释放资源：用try-finally确保，避免异常退出残留资源。
3. 监控显存：用NVIDIA Exporter+Prometheus+Grafana设阈值，超过告警，及时重启。
4. 版本选择：用稳定版本（如8.4+），旧版本有已知泄漏问题。

4.3 常见误区

• 误区1：Python GC会释放TensorRT对象→不会，必须手动del。
• 误区2：序列化后中间对象不用管→中间对象的显存和引擎独立，必须释放。
• 误区3：显存涨是模型大→模型显存固定，持续涨的肯定是泄漏。

五、文章总结

生产环境TensorRT显存泄漏，本质是C++核心对象和GPU显存的生命周期没正确管理。从序列化到推理的每个环节，都要检查资源是否释放，核心是：手动删除所有TensorRT核心对象，复用上下文和显存，服务退出时强制释放资源。定位时通过监控显存、拆分流程、查看日志，可快速找到问题，解决后大幅提升服务稳定性，避免OOM故障。