yjh0410 1 năm trước cách đây
mục cha
commit
ef5f0c6d51
3 tập tin đã thay đổi với 3 bổ sung0 xóa
  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