5个纯PyTorch/NumPy目标检测实现：YOLOv3 Tiny、Faster R-CNN RPN与TensorRT轻量化

简介本资源是一份面向深度学习初学者与实践者的Python目标检测代码实现教程聚焦YOLO、SSD、Faster R-CNN等主流算法在TensorFlow框架下的完整落地流程解决图像中物体识别与精确定位的核心问题。压缩包共308个文件涵盖210个Python源码含数据加载、模型构建、训练循环、评估与预测脚本、26个Protocol Buffer定义文件用于模型配置与标签映射、16个Markdown说明文档含环境配置、数据集格式解析与运行指南、15个config配置文件适配不同检测器结构及配套图片、Shell脚本与Jupyter Notebook示例整体4.05MB轻量易部署。已有3263人学习下载内容组织清晰Chapter_5目录结构严格对应理论章节——从数据预处理、模型选型、RPN与特征提取层实现到训练优化、指标评估及模型保存/推理全流程代码可直接运行辅以proto与config文件体现工业级配置规范是系统掌握TensorFlow目标检测工程实践的高价值参考材料。1. 这不是又一个YOLO复刻5个可直接跑通的目标检测Python实现覆盖单阶段/两阶段/轻量化全链路你手头刚拿到一张工业缺陷图想快速验证是否能检出微小划痕或者正带学生做课程设计需要避开PyTorch Lightning封装黑盒从dataloader、anchor生成、loss计算到NMS全部亲手过一遍——这时候翻遍GitHub要么是动辄3000行的完整训练框架含分布式、混合精度、模型服务器要么是只贴了3行model.predict()的“教学代码”。本资源包里这5个Python目标检测实现全部满足纯NumPyOpenCV或PyTorch原生API无第三方训练框架依赖每个实现都附带可直接运行的inference脚本和真实标注数据样例所有核心模块如IoU计算、anchor匹配、损失函数用不到50行代码写透逻辑。适合两类人一是想搞懂Faster R-CNN中RPN如何与backbone特征对齐的新手二是需要在嵌入式设备上裁剪检测头、必须看清每一层tensor shape变化的部署工程师。它不教你怎么调参但保证你改完第7行代码就能看到loss下降曲线。2. 从零构建单阶段检测器YOLOv3 Tiny的PyTorch原生实现与张量流解析YOLO系列之所以成为工业界首选并非仅因速度而是其将目标检测彻底解耦为“特征提取→网格预测→后处理”三段式流水线每段都可独立替换。本实现基于YOLOv3 Tiny结构但刻意剥离了Darknet53的C编译依赖全程使用PyTorch nn.Module定义重点暴露三个关键张量流节点输入图像经backbone输出的feature map shape、grid anchor映射后的pred tensor维度、以及NMS前原始bbox坐标的归一化逻辑。这种“裸写”方式牺牲了部分训练速度却让每个tensor的batch_size、channel、height、width维度变化完全可视——当你发现某次修改anchor尺寸后loss爆炸能立刻定位到是pred tensor的channel维度没同步更新。2.1 backbone与特征金字塔的张量对齐YOLOv3 Tiny采用两级下采样stride4和stride8对应两个检测头。我们不使用预训练权重而是从零初始化卷积层强制观察特征图尺寸收缩规律import torch import torch.nn as nn class YOLOv3TinyBackbone(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, 3, padding1) # 输入: [B,3,H,W] → 输出: [B,16,H,W] self.pool1 nn.MaxPool2d(2, 2) # → [B,16,H/2,W/2] self.conv2 nn.Conv2d(16, 32, 3, padding1) # → [B,32,H/2,W/2] self.pool2 nn.MaxPool2d(2, 2) # → [B,32,H/4,W/4] # 注意此处pool2后不再接conv直接作为第一个检测头输入 # 第二个检测头需从conv2输出分支经额外conv降维至16通道 self.conv_branch nn.Conv2d(32, 16, 1) # → [B,16,H/2,W/2] def forward(self, x): x torch.relu(self.conv1(x)) x self.pool1(x) x torch.relu(self.conv2(x)) feat_large self.pool2(x) # [B,32,H/4,W/4] → 第一个检测头输入 feat_small self.conv_branch(x) # [B,16,H/2,W/2] → 第二个检测头输入 return feat_large, feat_small参数说明feat_large尺寸为[B,32,H//4,W//4]对应13×13网格feat_small为[B,16,H//2,W//2]对应26×26网格。关键点在于两个特征图的H/W必须严格为整数倍关系此处26/132否则anchor映射时坐标会错位。若你替换为ResNet18需手动添加nn.AdaptiveAvgPool2d((13,13))确保输出尺寸一致而非依赖网络自动下采样。2.2 预测头与anchor匹配的tensor shape推演YOLO的预测头本质是1×1卷积将特征图每个像素点映射为多个anchor的偏移量。本实现定义2个anchor针对小目标和2个anchor针对大目标每个anchor预测5个值x,y,w,h,objectnessclass DetectionHead(nn.Module): def __init__(self, in_channels, num_anchors2): super().__init__() # 每个anchor预测5值tx,ty,tw,th,obj C类置信度此处C1二分类 self.conv nn.Conv2d(in_channels, num_anchors * 6, 1) def forward(self, x): pred self.conv(x) # 输入[B,C_in,H,W] → 输出[B,num_anchors*6,H,W] B, C, H, W pred.shape # 重排为[B, num_anchors, 6, H, W]便于后续操作 pred pred.view(B, -1, 6, H, W).permute(0,1,3,4,2) return pred # 形状: [B, A, H, W, 6] # 实例化两个头 head_large DetectionHead(32, num_anchors2) # 接feat_large head_small DetectionHead(16, num_anchors2) # 接feat_small关键逻辑pred.view(B,-1,6,H,W)中-1自动计算anchor数量此处为2permute将channel维度6移到末尾形成[B,A,H,W,6]结构。这样设计后后续计算IoU或loss时可直接对最后一维切片pred[...,0:2]取中心偏移pred[...,2:4]取宽高偏移。若你增加类别数如C3需将6改为5C且loss计算时需扩展one-hot标签维度。2.3 NMS前的坐标解码从归一化到绝对像素的三步转换YOLO输出的是相对于grid cell的归一化偏移需经三步解码才能得到图像坐标。本实现将每步tensor运算显式写出避免调用torchvision.ops.nms时丢失中间状态def decode_predictions(pred, anchors, stride, image_size): pred: [B,A,H,W,6], anchors: [[w1,h1],[w2,h2]], stride: 下采样倍数4或2 image_size: (H_img, W_img), 如(416,416) B, A, H, W, _ pred.shape # 步骤1生成grid坐标每个cell的左上角 grid_y, grid_x torch.meshgrid(torch.arange(H), torch.arange(W)) grid_x grid_x.unsqueeze(0).unsqueeze(0).float() # [1,1,H,W] grid_y grid_y.unsqueeze(0).unsqueeze(0).float() # 步骤2解码中心坐标sigmoid后乘stride加grid tx_ty torch.sigmoid(pred[..., 0:2]) # [B,A,H,W,2] bx_by tx_ty torch.stack([grid_x, grid_y], dim-1) # [B,A,H,W,2] bx_by bx_by * stride # 转为绝对坐标 # 步骤3解码宽高exp后乘anchor tw_th pred[..., 2:4] bw_bh torch.exp(tw_th) * torch.tensor(anchors).float().view(1,A,1,1,2) # 合并为[x1,y1,x2,y2]格式用于NMS x1 bx_by[..., 0] - bw_bh[..., 0] / 2 y1 bx_by[..., 1] - bw_bh[..., 1] / 2 x2 bx_by[..., 0] bw_bh[..., 0] / 2 y2 bx_by[..., 1] bw_bh[..., 1] / 2 boxes torch.stack([x1,y1,x2,y2], dim-1) # [B,A,H,W,4] scores torch.sigmoid(pred[..., 4]) * torch.sigmoid(pred[..., 5]) # obj * cls return boxes, scores # 使用示例 anchors_large [[116,90], [156,198]] # 大anchor适配13x13特征图 boxes_large, scores_large decode_predictions( pred_large, anchors_large, stride32, image_size(416,416) )参数说明stride32对应13×13特征图416/1332stride16对应26×26图。此处极易翻车若anchor尺寸与stride不匹配如用[116,90]配stride16解码后bbox会远超图像边界。实际项目中我们用k-means在训练集标注框上聚类得到anchor再按特征图stride分组分配。3. 两阶段检测器精简版Faster R-CNN的RPN模块手撕实现Faster R-CNN的精髓不在RCNN头而在RPNRegion Proposal Network如何用极小代价生成高质量候选框。本实现剥离了完整的Faster R-CNN仅保留RPN模块且用PyTorch原生API重写所有组件从feature map上的滑动窗口生成anchor到anchor与GT的IoU匹配再到ROI Pooling前的坐标变换。它不追求SOTA指标但保证你能看清为什么RPN的loss包含rpn_loss_cls和rpn_loss_bbox两项为什么正样本anchor要限制在图像内为什么ROI Pooling的输出尺寸必须固定3.1 Anchor生成在feature map上铺满可学习的先验框RPN不依赖预设anchor而是通过卷积核在feature map每个位置生成k个anchor。本实现用nn.Conv2d替代传统滑动窗口将anchor生成转化为卷积操作class AnchorGenerator: def __init__(self, base_size16, scales[8,16,32], ratios[0.5,1,2]): self.base_size base_size self.scales torch.tensor(scales).float() self.ratios torch.tensor(ratios).float() # 预计算所有anchor的宽高相对base_size self.anchor_sizes self.base_size * self.scales.unsqueeze(1) self.anchor_areas self.anchor_sizes ** 2 self.anchor_wh [] for area in self.anchor_areas: for ratio in self.ratios: w torch.sqrt(area / ratio) h area / w self.anchor_wh.append([w.item(), h.item()]) self.anchor_wh torch.tensor(self.anchor_wh) # [k,2], k9 def generate_anchors(self, feat_map_shape): feat_map_shape: (B,C,H,W) 返回: [H*W*k, 4] 的anchor坐标x1,y1,x2,y2归一化到[0,1] H, W feat_map_shape[2], feat_map_shape[3] # 生成feature map上每个点的坐标以pixel为单位 shift_x torch.arange(0, W * 16, 16) # stride16对应原图坐标 shift_y torch.arange(0, H * 16, 16) shift_x, shift_y torch.meshgrid(shift_x, shift_y) shifts torch.stack([shift_x, shift_y, shift_x, shift_y], dim-1) # [H,W,4] shifts shifts.view(-1, 4) # [H*W,4] # 每个shift位置生成k个anchor A self.anchor_wh.shape[0] K shifts.shape[0] anchors shifts.view(K, 1, 4) self.anchor_wh.view(1, A, 4) # [K,A,4] anchors anchors.view(-1, 4) # [K*A,4] # 裁剪到图像边界416x416 anchors[:, 0::2].clamp_(min0, max415) # x1,x2 anchors[:, 1::2].clamp_(min0, max415) # y1,y2 return anchors # 使用 anchor_gen AnchorGenerator() anchors anchor_gen.generate_anchors((1,512,26,26)) # [26*26*9,4] print(f生成{len(anchors)}个anchor)原理说明base_size16表示在feature map上每个点对应原图16×16区域scales和ratios生成9种anchor尺寸。关键约束所有anchor必须被裁剪到图像内clamp_否则后续计算IoU时会出现负面积导致梯度异常。这是RPN训练稳定性的第一道防线。3.2 Anchor与GT匹配IoU阈值与正负样本平衡策略RPN的loss计算依赖anchor与GT框的匹配结果。本实现不调用torchvision.ops.box_iou而是手写IoU计算并严格执行“正样本IoU0.7负样本IoU0.3中间样本忽略”的规则def compute_iou(boxes1, boxes2): boxes1: [N,4], boxes2: [M,4] → iou: [N,M] area1 (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1]) area2 (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1]) lt torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] rb torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] wh (rb - lt).clamp(min0) # [N,M,2] inter wh[:, :, 0] * wh[:, :, 1] # [N,M] iou inter / (area1[:, None] area2[None, :] - inter) return iou def match_anchors_to_gt(anchors, gt_boxes, pos_iou_thresh0.7, neg_iou_thresh0.3): anchors: [A,4], gt_boxes: [G,4] 返回: labels [A], bbox_targets [A,4] if len(gt_boxes) 0: return torch.zeros(len(anchors), dtypetorch.long), torch.zeros(len(anchors), 4) iou_matrix compute_iou(anchors, gt_boxes) # [A,G] max_iou, argmax_iou iou_matrix.max(dim1) # [A], [A] labels torch.full((len(anchors),), -1, dtypetorch.long) # -1: ignore labels[max_iou pos_iou_thresh] 1 # positive labels[max_iou neg_iou_thresh] 0 # negative # 确保每个GT至少有一个正样本避免漏检 gt_argmax_iou, _ iou_matrix.max(dim0) # [G] _, gt_argmax_idx torch.where(iou_matrix gt_argmax_iou) labels[gt_argmax_idx] 1 # 计算回归目标dx,dy,dw,dh gt_boxes_matched gt_boxes[argmax_iou] # [A,4] # 公式dx (gx - ax)/aw, dy (gy - ay)/ah, dw log(gw/aw), dh log(gh/ah) ax (anchors[:, 0] anchors[:, 2]) / 2 ay (anchors[:, 1] anchors[:, 3]) / 2 aw anchors[:, 2] - anchors[:, 0] ah anchors[:, 3] - anchors[:, 1] gx (gt_boxes_matched[:, 0] gt_boxes_matched[:, 2]) / 2 gy (gt_boxes_matched[:, 1] gt_boxes_matched[:, 3]) / 2 gw gt_boxes_matched[:, 2] - gt_boxes_matched[:, 0] gh gt_boxes_matched[:, 3] - gt_boxes_matched[:, 1] targets torch.stack([ (gx - ax) / aw, (gy - ay) / ah, torch.log(gw / aw 1e-8), torch.log(gh / ah 1e-8) ], dim1) return labels, targets # 示例匹配100个GT框 gt_boxes torch.tensor([[100,100,150,150],[200,200,250,250]]) # [2,4] labels, targets match_anchors_to_gt(anchors[:1000], gt_boxes) print(f正样本数: {(labels1).sum()}, 负样本数: {(labels0).sum()})血泪经验torch.log(gw/aw 1e-8)中的1e-8是后悔药——当gw或aw为0时防止log(0)崩溃。更隐蔽的坑是若GT框宽高为0标注错误会导致targets中出现inf进而使loss变为nan。我们在数据加载时强制过滤w5 or h5的GT框。3.3 RPN Loss计算分类与回归的梯度协同机制RPN的loss由两部分组成rpn_loss_clsanchor是否含物体和rpn_loss_bbox框坐标回归。本实现用FocalLoss替代原始交叉熵缓解正负样本极度不平衡问题并对回归loss加权正样本权重为10import torch.nn.functional as F def rpn_loss(rpn_cls_logits, rpn_bbox_pred, labels, bbox_targets): rpn_cls_logits: [B,A*2,H,W] → 二分类logitsobj/no-obj rpn_bbox_pred: [B,A*4,H,W] → 回归预测 labels: [A] → -1/0/1 bbox_targets: [A,4] B, C_cls, H, W rpn_cls_logits.shape A C_cls // 2 # 分类loss只计算label!-1的anchor cls_logits rpn_cls_logits.permute(0,2,3,1).reshape(B*H*W*A, 2) # [B*H*W*A,2] labels_flat labels.unsqueeze(0).expand(B,-1).reshape(-1) # [B*H*W*A] valid_mask labels_flat 0 cls_logits_valid cls_logits[valid_mask] labels_valid labels_flat[valid_mask] # Focal Lossalpha0.25, gamma2 ce_loss F.cross_entropy(cls_logits_valid, labels_valid, reductionnone) pt torch.exp(-ce_loss) focal_weight (1-pt)**2 * 0.25 rpn_loss_cls (focal_weight * ce_loss).mean() # 回归loss只计算正样本label1 bbox_pred rpn_bbox_pred.permute(0,2,3,1).reshape(B*H*W*A, 4) # [B*H*W*A,4] pos_mask labels_flat 1 if pos_mask.sum() 0: bbox_pred_pos bbox_pred[pos_mask] bbox_targets_pos bbox_targets.repeat(B,1)[pos_mask] # [N_pos,4] # Smooth L1 Loss loss_bbox F.smooth_l1_loss(bbox_pred_pos, bbox_targets_pos, reductionnone) rpn_loss_bbox loss_bbox.mean() * 10.0 # 正样本权重 else: rpn_loss_bbox torch.tensor(0.0) return rpn_loss_cls, rpn_loss_bbox # 模拟前向输出 rpn_cls_logits torch.randn(1, 18, 26, 26) # A9 → 9*218 rpn_bbox_pred torch.randn(1, 36, 26, 26) # 9*436 rpn_loss_cls, rpn_loss_bbox rpn_loss( rpn_cls_logits, rpn_bbox_pred, labels[:1000], targets[:1000] ) print(fRPN Loss: cls{rpn_loss_cls:.4f}, bbox{rpn_loss_bbox:.4f})避坑 / 常见问题 / 排查现象1训练初期rpn_loss_bbox持续为0原因pos_mask.sum() 0即没有anchor被标记为正样本。通常因anchor尺寸与GT框严重不匹配如GT平均宽高为20而anchor最小为32导致所有IoU0.3。解决检查AnchorGenerator中base_size和scales是否与数据集GT统计分布匹配打印max_iou直方图确认分布。现象2rpn_loss_cls震荡剧烈准确率在50%附近徘徊原因Focal Loss中gamma2放大了难样本权重但若正负样本比例超过1:1000仍需降低alpha如设为0.1。解决在rpn_loss中临时添加print(fpos:neg {pos_mask.sum()}:{(labels_flat0).sum()})若比例失衡则调整IoU阈值或增加anchor数量。现象3反向传播时报错RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation原因bbox_targets在match_anchors_to_gt中被原地修改如clamp_而该tensor被多次引用。解决所有涉及clamp_、fill_的操作前加.clone()如anchors[:, 0::2].clamp_(min0, max415)→anchors anchors.clone(); anchors[:, 0::2].clamp_(min0, max415)。现象4NMS后剩余框数极少5但GT有20个原因decode_predictions中stride参数错误如应为16却传入32导致解码坐标超出图像范围NMS时被过滤。解决在decode_predictions末尾添加assert (boxes0).all() and (boxes416).all()快速定位坐标溢出。4. 轻量化检测器实战MobileNetV2SSD的TensorRT部署准备当你的模型要跑在Jetson Nano或RK3399上参数量和FLOPs不再是唯一指标内存带宽占用、tensor shape对齐、op融合友好度才是瓶颈。本实现基于MobileNetV2 backbone SSD head但所有模块均按TensorRT优化要求重构禁用BN层替换为fused convscale、确保所有卷积stride为2的幂、将SSD的多尺度特征图concat操作改为逐层处理。更重要的是我们提供完整的ONNX导出脚本和shape inference验证避免部署时出现“dynamic shape not supported”错误。4.1 MobileNetV2 backbone的TensorRT友好改造标准MobileNetV2的Inverted Residual Block包含nn.BatchNorm2d而TensorRT对BN融合有严格要求需与前一conv紧邻。本实现将BN参数折叠进conv权重生成纯convrelu结构import torch.nn as nn import torch.nn.functional as F class InvertedResidualTRT(nn.Module): def __init__(self, inp, oup, stride, expand_ratio): super().__init__() hidden_dim int(round(inp * expand_ratio)) self.use_res_connect stride 1 and inp oup self.stride stride layers [] # 1. Pointwise conv (expand) if expand_ratio ! 1: layers.append(nn.Conv2d(inp, hidden_dim, 1, biasFalse)) layers.append(nn.ReLU6(inplaceTrue)) # 2. Depthwise conv layers.extend([ nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groupshidden_dim, biasFalse), nn.ReLU6(inplaceTrue) ]) # 3. Pointwise conv (project) layers.append(nn.Conv2d(hidden_dim, oup, 1, biasFalse)) self.conv nn.Sequential(*layers) # 关键预计算BN融合后的conv权重部署时无需BN self.fused_conv_weight None self.fused_conv_bias None def fuse_bn(self, bn_layer): 将bn_layer参数融合到最后一层conv conv self.conv[-1] w conv.weight bn_w bn_layer.weight bn_b bn_layer.bias bn_rm bn_layer.running_mean bn_rv bn_layer.running_var eps bn_layer.eps # 计算融合后权重w w * (bn_w / sqrt(bn_rv eps)) # 偏置b (bn_w / sqrt(...)) * (b - bn_rm) bn_b std torch.sqrt(bn_rv eps) fused_w w * (bn_w / std).view(-1, 1, 1, 1) if conv.bias is not None: fused_b (bn_w / std) * (conv.bias - bn_rm) bn_b else: fused_b bn_b - (bn_w / std) * bn_rm self.fused_conv_weight fused_w self.fused_conv_bias fused_b # 移除BN层后续forward只用fused conv self.conv nn.Sequential(*list(self.conv.children())[:-1]) def forward(self, x): if self.fused_conv_weight is not None: # 使用融合后conv x F.conv2d(x, self.fused_conv_weight, self.fused_conv_bias, strideself.stride, padding1, groups1) else: x self.conv(x) return x # 使用示例训练后调用fuse_bn block InvertedResidualTRT(32, 64, stride2, expand_ratio6) # 假设已训练好bn_layer为block.conv[-2]即Depthwise后的BN # block.fuse_bn(bn_layer)参数说明fuse_bn方法在训练完成后调用将BN参数永久写入conv权重。TensorRT部署时必须确保所有BN都被融合否则会报错Unsupported layer type: BatchNormalization。我们提供check_fused_model(model)函数遍历所有module检查是否存在nn.BatchNorm2d。4.2 SSD多尺度head的静态shape导出SSD在不同feature map上预测不同尺寸的box但TensorRT要求所有tensor shape在导出时确定。本实现强制指定各层输出尺寸禁用动态resizeclass SSDHeadTRT(nn.Module): def __init__(self, in_channels_list, num_classes2, num_anchors_list[4,6,6,4,4,4]): super().__init__() self.num_classes num_classes self.loc_layers nn.ModuleList() self.cls_layers nn.ModuleList() for i, (in_c, na) in enumerate(zip(in_channels_list, num_anchors_list)): # 定制化输出尺寸layer0→38x38, layer1→19x19, layer2→10x10... self.loc_layers.append(nn.Conv2d(in_c, na*4, 3, padding1)) self.cls_layers.append(nn.Conv2d(in_c, na*num_classes, 3, padding1)) # 预定义各层输出尺寸TensorRT必需 self.output_shapes [ (38,38), (19,19), (10,10), (5,5), (3,3), (1,1) ] def forward(self, features): features: list of [B,C,H,W] tensors 返回: loc_preds [B, total_anchors, 4], cls_preds [B, total_anchors, C] loc_preds [] cls_preds [] for i, (feat, loc_layer, cls_layer) in enumerate( zip(features, self.loc_layers, self.cls_layers) ): # 强制校验feature shape与预设一致 assert feat.shape[2:] self.output_shapes[i], \ fFeature {i} shape {feat.shape[2:]} ! expected {self.output_shapes[i]} loc_pred loc_layer(feat) # [B,na*4,H,W] cls_pred cls_layer(feat) # [B,na*C,H,W] # 展平为[B, H*W*na, 4]和[B, H*W*na, C] B, _, H, W loc_pred.shape na loc_pred.shape[1] // 4 loc_pred loc_pred.permute(0,2,3,1).reshape(B, -1, 4) cls_pred cls_pred.permute(0,2,3,1).reshape(B, -1, self.num_classes) loc_preds.append(loc_pred) cls_preds.append(cls_pred) # concat所有尺度 loc_preds torch.cat(loc_preds, dim1) # [B, total_anchors, 4] cls_preds torch.cat(cls_preds, dim1) # [B, total_anchors, C] return loc_preds, cls_preds # 构建特征图列表模拟MobileNetV2输出 features [ torch.randn(1, 24, 38, 38), # layer0 torch.randn(1, 32, 19, 19), # layer1 torch.randn(1, 96, 10, 10), # layer2 torch.randn(1, 1280, 5, 5), # layer3 ] head SSDHeadTRT([24,32,96,1280]) loc_out, cls_out head(features) print(fSSD输出: loc{loc_out.shape}, cls{cls_out.shape})关键逻辑assert feat.shape[2:] self.output_shapes[i]在forward中强制校验确保ONNX导出时不会出现dynamic shape。若你替换backbone必须同步更新output_shapes和in_channels_list否则导出失败。4.3 ONNX导出与TensorRT引擎生成脚本导出ONNX时必须指定dynamic_axes为空并用--opset-version 11兼容TensorRT 7# 导出命令在Python脚本中调用 torch.onnx.export( model, dummy_input, ssd_mobilenetv2.onnx, input_names[input], output_names[loc, cls], dynamic_axes{}, # 关键禁用dynamic axes opset_version11, verboseFalse )# 验证ONNX shape部署前必做 import onnx model onnx.load(ssd_mobilenetv2.onnx) for node in model.graph.node: if node.op_type Conv: # 检查conv的kernel_shape是否为[3,3] kernel [attr.ints for attr in node.attribute if attr.namekernel_shape][0] assert kernel [3,3], fConv kernel {kernel} not supported # TensorRT引擎生成Python API import tensorrt as trt TRT_LOGGER trt.Logger(trt.Logger.WARNING) builder trt.Builder(TRT_LOGGER) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, TRT_LOGGER) with open(ssd_mobilenetv2.onnx, rb) as f: parser.parse(f.read()) engine builder.build_cuda_engine(network) # 保存engine with open(ssd.trt, wb) as f: f.write(engine.serialize())避坑 / 常见问题 / 排查现象1ONNX导出时报错Exporting aten::size operator is not supported原因模型中存在x.size(0)等动态shape操作。解决全部替换为x.shape[0]并在forward开头添加B x.shape[0]后续用B代替。现象2TensorRT构建引擎时卡死或内存溢出原因builder.max_workspace_size默认为1GB而SSD head参数量大。解决设置builder.max_workspace_size 1 301GB或更高如1本文还有配套的精品资源点击获取