PyTorch转TensorRT引擎后加载目标检测模型报错，PyTorch格式正常

本人编程能力有限，尝试将目标检测模型通过TensorRT加载到GPU以提升运行速度时遇到问题。脚本在PyTorch格式下配合CUDA可正常运行，但把PyTorch模型导出为TensorRT的.engine格式后运行出现如下错误：

[04/05/2023-19:41:54] [TRT] [I] [MemUsageChange] Init CUDA: CPU +499, GPU +0, now: CPU 11536, GPU 1240 (MiB)
[04/05/2023-19:41:54] [TRT] [I] Loaded engine size: 5 MiB
[04/05/2023-19:41:55] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +1186, GPU +408, now: CPU 12783, GPU 1654 (MiB)
[04/05/2023-19:41:55] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in engine deserialization: CPU +0, GPU +0, now: CPU 0, GPU 0 (MiB)
[04/05/2023-19:41:55] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +0, GPU +8, now: CPU 12776, GPU 1654 (MiB)
[04/05/2023-19:41:55] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +0, now: CPU 0, GPU 0 (MiB)
Adding AutoShape... 
(1376, 576, 2064, 864)
Traceback (most recent call last):
  File "c:\Users\HALLÅÅ\Documents\CHEATS\scripts\main copy.py", line 150, in <module>
    df= model(screenshot, size=640).pandas().xyxy[0]
  File "C:\Users\HALLÅÅ\Documents\CHEATS\cheats-env\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\HALLÅÅ\Documents\CHEATS\cheats-env\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "C:\Users\HALLÅÅ\Documents\CHEATS\yolov5\models\common.py", line 704, in forward
    y = self.model(x, augment=augment)  # forward
  File "C:\Users\HALLÅÅ\Documents\CHEATS\cheats-env\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\HALLÅÅ\Documents\CHEATS\yolov5\models\common.py", line 536, in forward
    assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
AssertionError: input size torch.Size([1, 3, 288, 640]) not equal to max model size (1, 3, 640, 640)

问题分析 #

报错核心是输入尺寸不匹配：输入图像经过预处理后尺寸为[1,3,288,640]，但TensorRT引擎要求的固定输入尺寸是[1,3,640,640]。原因是导出TensorRT引擎时未开启动态尺寸支持，引擎只能处理固定的640x640输入；而输入图像宽高比非1:1，AutoShape模块按比例缩放后未补全到640x640，触发断言错误。

解决办法 #

• 开启动态尺寸导出引擎：导出模型时添加--dynamic参数，让引擎支持不同宽高比的输入：

  python export.py --weights your_model.pt --include engine --dynamic
  

  若用Python代码构建引擎，需手动指定动态维度范围，比如设置min_shape=(1,3,320,320)、opt_shape=(1,3,640,640)、max_shape=(1,3,1280,1280)。

• 强制输入尺寸与引擎一致：如果不需要动态尺寸，运行时强制图像缩放到640x640：
  调用模型时设置rect=False，关闭按比例缩放，直接拉伸图像到指定尺寸：

  df = model(screenshot, size=640, rect=False).pandas().xyxy[0]
  

  也可以手动对图像做resize和padding处理，确保输入尺寸严格为640x640。

• 临时禁用断言（不推荐）：直接注释掉yolov5/models/common.py第536行的断言代码，但可能导致后续推理出现未知错误，仅适合临时测试场景。

内容的提问来源于stack exchange，提问作者Klurig Gubbe