DETR模型转TensorRT时如何指定部分层采用FP16混合精度推理？

方案1：ONNX层标记法（适配你的PyTorch->ONNX->TensorRT流程） #

在导出的ONNX模型中给精度敏感层添加FP32强制属性，TensorRT转换时会自动识别：

1. 用onnx库加载DETR的ONNX模型：

   import onnx
   model = onnx.load("detr_base.onnx")
   

2. 遍历模型节点，给需要保留FP32的节点（比如DETR的分类头、回归头、注意力归一化层）添加精度属性：

   # 先通过Netron工具查看目标节点名称，替换为实际匹配规则
   target_node_keywords = ["class_head", "bbox_head", "layer_norm"]
   for node in model.graph.node:
       if any(keyword in node.name for keyword in target_node_keywords):
           attr = onnx.helper.make_attribute("precision", "FP32")
           node.attribute.append(attr)
   onnx.save(model, "detr_marked_fp32.onnx")
   

3. 用TensorRT转换时开启FP16模式，工具会跳过标记为FP32的节点：

   trtexec --onnx=detr_marked_fp32.onnx --saveEngine=detr_mixed.engine --fp16
   

方案2：TensorRT API手动指定精度 #

直接通过TensorRT Python API构建引擎时，强制指定目标层的精度：

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
parser.parse_from_file("detr_base.onnx")

# 遍历所有层，设置FP32约束
target_layer_keywords = ["class_head", "bbox_head", "layer_norm"]
for layer_idx in range(network.num_layers):
    layer = network.get_layer(layer_idx)
    if any(keyword in layer.name for keyword in target_layer_keywords):
        # 强制层精度为FP32
        layer.precision = trt.DataType.FLOAT
        # 同步输入输出张量的精度
        for in_idx in range(layer.num_inputs):
            network.get_tensor(layer.get_input(in_idx)).dtype = trt.DataType.FLOAT
        for out_idx in range(layer.num_outputs):
            network.get_tensor(layer.get_output(out_idx)).dtype = trt.DataType.FLOAT

# 开启FP16模式构建引擎
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)
engine = builder.build_serialized_network(network, config)

# 保存引擎
with open("detr_mixed.engine", "wb") as f:
    f.write(engine)

核心注意事项 #

• 先用Netron可视化DETR模型，精准定位对精度敏感的层（通常是最后输出的分类/回归头、注意力模块中的LayerNorm、query/key投影层）
• 从最关键的输出头开始测试，逐步扩大FP32层范围，平衡精度与性能
• 避免对大量层设置FP32，否则会丧失混合精度的加速优势

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