yjh0410 1 жил өмнө
parent
commit
ef5f0c6d51

+ 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