Jetson Orin Nano上PyTorch转ONNX再转TensorRT引擎失败求助

问题背景 #

在Jetson Orin Nano设备上转换UNICOM仓库中的ViT-B/32模型时，遇到两个核心问题：

1. 使用torch.onnx.dynamo_export导出ONNX后，通过trtexec转换为TensorRT引擎时，提示找不到unicom_vision_transformer_PatchEmbedding_patch_embed_1插件。
2. 改用torch.onnx.export导出ONNX时，触发RuntimeError: _Map_base::at错误。

相关代码与信息 #

初始导出ONNX代码（dynamo_export版本） #

import torch
import onnx
import onnxruntime

from unicom.vision_transformer import build_model

if __name__ == '__main__':
    model_name = "ViT-B/32"
    model_name_fp16 = "FP16-ViT-B-32"
    onnx_model_path = f"{model_name_fp16}.onnx"

    model = build_model(model_name)
    model.eval()
    model = model.to('cuda')
    torch_input = torch.randn(1, 3, 224, 224).to('cuda')

    onnx_program = torch.onnx.dynamo_export(model, torch_input)
    onnx_program.save(onnx_model_path)

    onnx_model = onnx.load(onnx_model_path)
    onnx.checker.check_model(onnx_model_path)

TensorRT转换命令 #

/usr/src/tensorrt/bin/trtexec --onnx=FP16-ViT-B-32.onnx --saveEngine=FP16-ViT-B-32.trt --workspace=1024 --fp16

关键错误信息（dynamo_export导出后） #

No importer registered for op: unicom_vision_transformer_PatchEmbedding_patch_embed_1. Attempting to import as plugin.
3: getPluginCreator could not find plugin: unicom_vision_transformer_PatchEmbedding_patch_embed_1 version: 1

PatchEmbedding类代码 #

class PatchEmbedding(nn.Module):
    def __init__(self, input_size=224, patch_size=32, in_channels: int = 3, dim: int = 768):
        super().__init__()
        if isinstance(input_size, int):
            input_size = (input_size, input_size)
        if isinstance(patch_size, int):
            patch_size = (patch_size, patch_size)
        H = input_size[0] // patch_size[0]
        W = input_size[1] // patch_size[1]
        self.num_patches = H * W
        self.proj = nn.Conv2d(
            in_channels, dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x

环境信息 #

• TensorRT版本：8.6.2.3
• GPU：Jetson Orin Nano
• CUDA版本：12.2
• PyTorch版本：2.3.0
• ONNX版本：1.16.1

改用torch.onnx.export的错误 #

RuntimeError: _Map_base::at

解决方案 #

问题1：dynamo_export生成自定义OP导致TensorRT无法识别 #

原因：torch.onnx.dynamo_export默认会将模型中的子模块（如PatchEmbedding）融合为自定义OP，而非拆解为标准ONNX算子，TensorRT没有对应的导入器或插件，因此报错。

解决方法：
禁用dynamo的算子融合功能，强制导出标准ONNX算子：

import torch
import onnx
from unicom.vision_transformer import build_model

if __name__ == '__main__':
    model_name = "ViT-B/32"
    model_name_fp16 = "FP16-ViT-B-32"
    onnx_model_path = f"{model_name_fp16}.onnx"

    model = build_model(model_name)
    model.eval()
    model = model.to('cuda')
    torch_input = torch.randn(1, 3, 224, 224).to('cuda')

    # 禁用算子融合，避免生成自定义OP
    export_options = torch.onnx.ExportOptions(
        fuse_ops=False,
        decomposition_table=torch.onnx.registry.get_decomposition_table()
    )
    onnx_program = torch.onnx.dynamo_export(model, torch_input, export_options=export_options)
    onnx_program.save(onnx_model_path)

    onnx_model = onnx.load(onnx_model_path)
    onnx.checker.check_model(onnx_model)

问题2：torch.onnx.export触发RuntimeError: _Map_base::at #

原因：该错误通常是由于导出时未明确指定输入维度、模型存在动态维度未配置，或未禁用梯度计算导致的。

解决方法：
使用torch.onnx.export时，添加必要配置并确保模型处于正确状态：

import torch
import onnx
from unicom.vision_transformer import build_model

if __name__ == '__main__':
    model_name = "ViT-B/32"
    model_name_fp16 = "FP16-ViT-B-32"
    onnx_model_path = f"{model_name_fp16}.onnx"

    model = build_model(model_name)
    model.eval()
    model = model.to('cuda')
    # 固定输入形状，避免动态维度问题
    torch_input = torch.randn(1, 3, 224, 224).to('cuda')

    # 用torch.no_grad()禁用梯度计算，避免导出时的张量问题
    with torch.no_grad():
        torch.onnx.export(
            model,
            torch_input,
            onnx_model_path,
            opset_version=17,  # TensorRT 8.6支持到opset 17，匹配版本
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes=None,  # 固定输入输出形状，如需动态可配置对应维度
            export_params=True,
            verbose=False
        )

    onnx_model = onnx.load(onnx_model_path)
    onnx.checker.check_model(onnx_model)

额外注意事项 #

1. TensorRT版本匹配：确保导出ONNX时的opset_version与TensorRT支持的版本一致（TensorRT 8.6支持opset 1-17）。
2. FP16转换：如果需要FP16引擎，导出时可以将模型转换为FP16：

   model = model.half()
   torch_input = torch_input.half()
   

   再配合trtexec的--fp16参数，确保转换顺利。
3. Jetson环境优化：Jetson设备上建议使用TensorRT官方提供的容器或预编译包，避免版本不兼容问题。

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