yjh0410 1 年之前
父节点
当前提交
ef5f0c6d51
共有 3 个文件被更改,包括 3 次插入0 次删除
  1. 1 0
      models/yolov3/yolov3.py
  2. 1 0
      models/yolov5/yolov5.py
  3. 1 0
      models/yolox/yolox.py

+ 1 - 0
models/yolov3/yolov3.py

@@ -35,6 +35,7 @@ class Yolov3(nn.Module):
         self.pyramid_feat_dims = self.backbone.feat_dims[-3:]
         ## Neck: SPP
         self.neck     = SPPF(cfg, self.pyramid_feat_dims[-1], self.pyramid_feat_dims[-1])
+        self.pyramid_feat_dims[-1] = self.neck.out_dim
         ## Neck: FPN
         self.fpn      = Yolov3FPN(cfg, self.pyramid_feat_dims)
         ## Head

+ 1 - 0
models/yolov5/yolov5.py

@@ -35,6 +35,7 @@ class Yolov5(nn.Module):
         self.pyramid_feat_dims = self.backbone.feat_dims[-3:]
         ## Neck: SPP
         self.neck     = SPPF(cfg, self.pyramid_feat_dims[-1], self.pyramid_feat_dims[-1])
+        self.pyramid_feat_dims[-1] = self.neck.out_dim
         ## Neck: FPN
         self.fpn      = Yolov5PaFPN(cfg, self.pyramid_feat_dims)
         ## Head

+ 1 - 0
models/yolox/yolox.py

@@ -35,6 +35,7 @@ class Yolox(nn.Module):
         self.pyramid_feat_dims = self.backbone.feat_dims[-3:]
         ## Neck: SPP
         self.neck     = SPPF(cfg, self.pyramid_feat_dims[-1], self.pyramid_feat_dims[-1])
+        self.pyramid_feat_dims[-1] = self.neck.out_dim
         ## Neck: FPN
         self.fpn      = YoloxPaFPN(cfg, self.pyramid_feat_dims)
         ## Head