Update nets/yolo.py

nets/yolo.py

@@ -1,20 +1,9 @@
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
-from nets.CSPdarknet import darknet
 from nets.CSPdarknet_tiny import darknet_tiny
-from nets.mobilenetv2 import mobilenet_v2
-from nets.shufflenet_v2 import shufflenet_v2
-from nets.ghostnet import ghostnet
-from nets.attention import cbam_block, eca_block, se_block, CA_Block
-attention_block = [se_block, cbam_block, eca_block, CA_Block]
 
 
-
-#-------------------------------------------------#
-#   卷积块 -> 卷积 + 标准化 + 激活函数
-#   Conv2d + BatchNormalization + LeakyReLU
-#-------------------------------------------------#
 class BasicConv(nn.Module):
     def __init__(self, in_channels, out_channels, kernel_size, stride=1):
         super(BasicConv, self).__init__()
@@ -29,9 +18,6 @@ class BasicConv(nn.Module):
         x = self.activation(x)
         return x
 
-#---------------------------------------------------#
-#   卷积 + 上采样
-#---------------------------------------------------#
 class Upsample(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(Upsample, self).__init__()
@@ -45,9 +31,6 @@ class Upsample(nn.Module):
         x = self.upsample(x)
         return x
 
-#---------------------------------------------------#
-#   最后获得yolov4的输出
-#---------------------------------------------------#
 def yolo_head(filters_list, in_filters):
     m = nn.Sequential(
         BasicConv(in_filters, filters_list[0], 3),
@@ -128,26 +111,9 @@ class CARAFE(nn.Module):
 #   yolo_body--MSFNet
 #---------------------------------------------------#
 class YoloBody(nn.Module):
-    def __init__(self, anchors_mask, num_classes, phi=0, backbone ='', pretrained=False):
+    def __init__(self, anchors_mask, num_classes, phi=0, backbone ='tiny', pretrained=False):
         super(YoloBody, self).__init__()
-
-        self.phi            = phi
-
-        if backbone == 'cspdarknet':
-            self.backbone  = darknet(pretrained)
-            self.conv_for_P5    = BasicConv(512,256,1)
-            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
-            self.upsample_1       = Upsample(256,128)
-            self.conv1 = BasicConv(256,128,1)
-            self.upsample_2 = CARAFE(128)
-            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],384)
-
-            if 1 <= self.phi and self.phi <= 4:
-                self.feat1_att      = attention_block[self.phi - 1](256)
-                self.feat2_att      = attention_block[self.phi - 1](512)
-                self.upsample_att   = attention_block[self.phi - 1](128)  
-        
-        elif backbone == 'tiny':
+        if backbone == 'tiny':
             self.backbone  = darknet_tiny(pretrained)
             self.conv_for_P5    = BasicConv(512,256,1)
             self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
@@ -156,67 +122,9 @@ class YoloBody(nn.Module):
             self.upsample_2 = CARAFE(128)
             self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],384)
 
-            if 1 <= self.phi and self.phi <= 4:
-                self.feat1_att      = attention_block[self.phi - 1](256)
-                self.feat2_att      = attention_block[self.phi - 1](512)
-                self.upsample_att   = attention_block[self.phi - 1](128)  
-
-        elif backbone == 'mobilenetv2':
-            self.backbone  = mobilenet_v2(pretrained)
-            self.conv_for_P5    = BasicConv(320,256,1)
-            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
-            self.upsample_1       = Upsample(256,128)
-            self.conv1 = BasicConv(256,128,1)
-            self.upsample_2 = CARAFE(128)
-            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],224)
-
-            if 1 <= self.phi and self.phi <= 4:
-                self.feat1_att      = attention_block[self.phi - 1](256)
-                self.feat2_att      = attention_block[self.phi - 1](512)
-                self.upsample_att   = attention_block[self.phi - 1](128)  
-
-        elif backbone == 'shufflenetv2':
-            self.backbone  = shufflenet_v2()
-            self.conv_for_P5    = BasicConv(1024,256,1)
-            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
-            self.upsample_1       = Upsample(256,128)
-            self.conv1 = BasicConv(256,128,1)
-            self.upsample_2 = CARAFE(128)
-            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],592)
-
-            if 1 <= self.phi and self.phi <= 4:
-                self.feat1_att      = attention_block[self.phi - 1](256)
-                self.feat2_att      = attention_block[self.phi - 1](512)
-                self.upsample_att   = attention_block[self.phi - 1](128)  
-
-        elif backbone == 'ghostnet':
-            self.backbone  = ghostnet()
-            self.conv_for_P5    = BasicConv(160,256,1)
-            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
-            self.upsample_1       = Upsample(256,128)
-            self.conv1 = BasicConv(256,128,1)
-            self.upsample_2 = CARAFE(128)
-            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],240)
-
-            if 1 <= self.phi and self.phi <= 4:
-                self.feat1_att      = attention_block[self.phi - 1](256)
-                self.feat2_att      = attention_block[self.phi - 1](512)
-                self.upsample_att   = attention_block[self.phi - 1](128)  
-
-        
-
     def forward(self, x):
 
-        #---------------------------------------------------#
-        #   生成CSPdarknet53_tiny的主干模型
-        #   feat1的shape为26,26,256
-        #   feat2的shape为13,13,512
-        #---------------------------------------------------#
         feat1, feat2 = self.backbone(x)
-        if 1 <= self.phi and self.phi <= 4:
-            feat1 = self.feat1_att(feat1)
-            feat2 = self.feat2_att(feat2)
-
         # 13,13,512 -> 13,13,256
         P5 = self.conv_for_P5(feat2)
         # 13,13,256 -> 13,13,512 -> 13,13,255
@@ -225,7 +133,7 @@ class YoloBody(nn.Module):
         P6 = self.conv_for_P5(feat2)
         P6_Upsample = self.upsample_1(P6)
 
-        # 13,13,256 -> 13,13,128 -> 26,26,128  卷积+轻量级上采样
+        # 13,13,256 -> 13,13,128 -> 26,26,128  
         P5 = self.conv1(P5)
         P5_Upsample = self.upsample_2(P5)