清除python TensorRT build之后常驻的显存

业务需要,需要TensorRT 模型转换的进程常驻,使用中发现每次build后会占据一个显存缓存,del也无法清除,后来找到了:
python - How to run TensorRT in multiple threads? - Stack Overflow

import pycuda.driver as cuda
from pycuda.tools import make_default_context
print("Create TensorRT Plan.")
cuda.init()  # Initialize CUDA
ctx = make_default_context()
TRT_LOGGER = trt.Logger()
build_result = self.build_engine(TRT_LOGGER)
ctx.pop()
del ctx

你可能感兴趣的:(清除python TensorRT build之后常驻的显存)