yjh0410 hai 1 ano
pai
achega
ef5f0c6d51
Modificáronse 3 ficheiros con 3 adicións e 0 borrados
  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