TRTWrapper导致ONNX模型超量显存占用

[lck1201]

代码如下

#from mmcv.tensorrt import TRTWrapper

import ipdb; ipdb.set_trace()

my_onnx_model = MyONNXModel(**model_config) #实例化，在GPU加载模型

问题描述：
正常而言，加载onnx模型耗费580MB显存，但加上 from mmcv.tensorrt import TRTWrapper，显存占用激增至1500+MB，发生在 my_onnx_model = MyONNXModel(**model_config) 这行代码。
尝试进入TRTWrapper查看，但没找到影响显存的代码，求助。

[lck1201]

补充：因为同时用到了ONNX模型与TensorRT模型，因此 import TRTWrapper

[lck1201]

补充：似乎是 TRTWrapper中import torch导致了这个问题

[grimoire]

Hi,
请问显存增加的时候是否加载了TensorRT模型？如果加载了，那么模型的min/max shape还有max workspace size分别是多少？
import torch导致的问题就不太清楚了，可能需要更多线索
如果方便的话麻烦提供更具体的复现脚本

[lck1201]

hi，显存增加的时候，仅仅加载了ONNX模型，也就是my_onnx_model = MyONNXModel(**model_config)这行代码。
具体什么底层原因不清楚，猜测是import torch时内部设了一些cuda参数。

在类MyONNXModel里，用这个函数加载onnx模型：onnxruntime.InferenceSession(model_path, sess_opt)

用TRTWrapper加载模型、推理，暂时没发现异常。

from mmcv.tensorrt import TRTWrapper

import ipdb; ipdb.set_trace()

my_onnx_model = MyONNXModel(**model_config) #实例化，在GPU加载模型

[zengjie617789]

mmcv 使用算子倒没有显存飙升，但是使用不使用mmcv中的封装，直接使用tensorrt的算子，倒是会显存增加。初步估计是mmcv对算子插件进行了优化。