YOLOv5系列（3）——YOLOv5修改网络结构

傷城~ 2022-11-06 02:53 501阅读 0赞

### 文章目录 ###

*  一、设置网络结构为mobilenet-V2
 *  二、添加注意力模块

# 一、设置网络结构为mobilenet-V2 #

首先，需要在models/common.py里，实现MobileNetv2的 bottleneck 和 Pwconv。  
1、Mobilenetv2的bottleneck: InvertedResidual

#mobilenet  Bottleneck  InvertedResidual  
    class BottleneckMOB(nn.Module):  
        #c1:inp  c2:oup s:stride  expand_ratio:t  
        def __init__(self, c1, c2, s, expand_ratio):  
            super(BottleneckMOB, self).__init__()  
            self.s = s  
            hidden_dim = round(c1 * expand_ratio)  
            self.use_res_connect = self.s == 1 and c1 == c2  
            if expand_ratio == 1:  
                self.conv = nn.Sequential(  
                    # dw  
                    nn.Conv2d(hidden_dim, hidden_dim, 3, s, 1, groups=hidden_dim, bias=False),  
                    nn.BatchNorm2d(hidden_dim),  
                    nn.ReLU6(inplace=True),  
                    # pw-linear  
                    nn.Conv2d(hidden_dim, c2, 1, 1, 0, bias=False),  
                    nn.BatchNorm2d(c2),  
                )  
            else:  
                self.conv = nn.Sequential(  
                    # pw  
                    nn.Conv2d(c1, hidden_dim, 1, 1, 0, bias=False),  
                    nn.BatchNorm2d(hidden_dim),  
                    nn.ReLU6(inplace=True),  
                    # dw  
                    nn.Conv2d(hidden_dim, hidden_dim, 3, s, 1, groups=hidden_dim, bias=False),  
                    nn.BatchNorm2d(hidden_dim),  
                    nn.ReLU6(inplace=True),  
                    # pw-linear  
                    nn.Conv2d(hidden_dim, c2, 1, 1, 0, bias=False),  
                    nn.BatchNorm2d(c2),  
                )  
    
        def forward(self, x):  
            if self.use_res_connect:  
                return x + self.conv(x)  
            else:  
                return self.conv(x)

2、Pointwise Convolution

class PW_Conv(nn.Module):  
        def __init__(self, c1, c2):  # ch_in, ch_out  
            super(PW_Conv, self).__init__()  
            self.conv = nn.Conv2d(c1, c2, 1, 1, 0, bias=False)  
            self.bn = nn.BatchNorm2d(c2)  
            self.act = nn.ReLU6(inplace=True)  
    
        def forward(self, x):  
            return self.act(self.bn(self.conv(x)))

接着需要在yolov5的读取模型配置文件的代码（models/yolo.py的parse\_model函数）进行修改，使得能够调用到上面的模块，只需修改下面这部分代码：

n = max(round(n * gd), 1) if n > 1 else n  # depth gain  
    if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, PW_Conv, BottleneckMOB]:  
        c1, c2 = ch[f], args[0]

并且需要在import引用处加入PW\_Conv,BottleneckMOB这两个模块：

from models.common import Conv, Bottleneck,SPP, DWConv, Focus, BottleneckCSP, Concat, NMS, autoShape, PW_Conv,BottleneckMOB

然后就是搭建我们的模型配置文件，我在yolov5s.yaml的基础上进行修改，将yolov5s的backbone替换成mobilenetv2，重新建立了一个模型配置文件yolov5-mobilenetV2.yaml：

# parameters  
    nc: 1  # number of classes  
    depth_multiple: 0.33  # model depth multiple  
    width_multiple: 0.50  # layer channel multiple  
    
    # anchors  
    anchors:  
      - [116,90, 156,198, 373,326]  # P5/32  
      - [30,61, 62,45, 59,119]  # P4/16  
      - [10,13, 16,30, 33,23]  # P3/8  
    
    # YOLOv5 backbone: mobilenet v2  
    backbone:  
      # [from, number, module, args]  
      [[-1, 1, nn.Conv2d, [32, 3, 2]],  # 0-P1/2   oup, k, s     640  
       [-1, 1, BottleneckMOB, [16, 1, 1]],  # 1-P2/4   oup, s, t 320  
       [-1, 2, BottleneckMOB, [24, 2, 6]],  #                    320  
       [-1, 1, PW_Conv, [256]],  #4  output p3                   160  
       [-1, 3, BottleneckMOB, [32, 2, 6]],  # 3-P3/8             160  
       [-1, 4, BottleneckMOB, [64, 1, 6]],  # 5                  80  
       [-1, 1, PW_Conv, [512]],  #7 output p4  6                 40  
       [-1, 3, BottleneckMOB, [96, 2, 6]],  # 7                  80  
       [-1, 3, BottleneckMOB, [160, 1, 6,]], #                   40  
       [-1, 1, BottleneckMOB, [320, 1, 6,]], #                   40  
       [-1, 1, nn.Conv2d, [1280, 1, 1]],     #                   40  
       [-1, 1, SPP, [1024, [5, 9, 13]]],  #11     #              40  
      ]  
    
    # YOLOv5 head  
    head:  
      [[-1, 3, BottleneckCSP, [1024, False]],  # 12             40  
    
       [-1, 1, Conv, [512, 1, 1]],                      #       40  
       [-1, 1, nn.Upsample, [None, 2, 'nearest']],      #       40  
       [[-1, 6], 1, Concat, [1]],  # cat backbone P4-7  #       80  
       [-1, 3, BottleneckCSP, [512, False]],  # 16      #       80  
    
       [-1, 1, Conv, [256, 1, 1]],                      #       80  
       [-1, 1, nn.Upsample, [None, 2, 'nearest']],      #       160  
       [[-1, 3], 1, Concat, [1]],  # cat backbone P3-4          160  
       [-1, 3, BottleneckCSP, [256, False]],            #       160  
       [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]],  # 21 (P3/8-small)   #        160  
    
       [-2, 1, Conv, [256, 3, 2]],                     #       160  
       [[-1, 17], 1, Concat, [1]],  # cat head P4      #       160  
       [-1, 3, BottleneckCSP, [512, False]],           #       160  
       [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]],  # 25 (P4/16-medium)  #       160  
    
       [-2, 1, Conv, [512, 3, 2]],                     #       160  
       [[-1, 13], 1, Concat, [1]],  # cat head P5-13   #      160  
       [-1, 3, BottleneckCSP, [1024, False]],          #      160  
       [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]],  # 29 (P5/32-large)           160  
    
       [[21, 25, 29], 1, Detect, [nc, anchors]],  # Detect(P5, P4, P3)     nc:number class, na:number of anchors  
      ]

到这我们就实现了将yolov5的backbone替换成了mobilenetv2。在使用时只需要将网络结构配置参数—cfg修改成 –cfg yolov5-mobilenet.yaml。  
训练指令：

python train.py --data coco.yaml --cfg yolov5-mobilenet.yaml--weights '' --batch-size 64

# 二、添加注意力模块 #

配置文件yolov5x\_se.yaml

# parameters
    nc: 15 # number of classes
    depth_multiple: 1 # model depth multiple
    width_multiple: 1 # layer channel multiple
    
    # anchors
    anchors:
      - [10, 13, 16, 30, 33, 23] # P3/8
      - [30, 61, 62, 45, 59, 119] # P4/16
      - [116, 90, 156, 198, 373, 326] # P5/32
    
    # YOLOv5 backbone
    backbone:
      # [from, number, module, args]
      [
        [-1, 1, Focus, [64, 3]], # 0-P1/2                 #1
        [-1, 1, Conv, [128, 3, 2]], # 1-P2/4              #2
        [-1, 3, C3, [128]], #3
        [-1, 1, Conv, [256, 3, 2]], # 3-P3/8              #4
        [-1, 9, C3, [256]], #5
        [-1, 1, Conv, [512, 3, 2]], # 5-P4/16             #6
        [-1, 9, C3, [512]], #7
        [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32            #8
        [-1, 1, SPP, [1024, [5, 9, 13]]], #9
        [-1, 3, C3, [1024, False]], # 9                   #10
        [-1, 1, SELayer, [1024, 4]], #10
      ]
    
    
    
    # YOLOv5 head
    head: [
        [-1, 1, Conv, [512, 1, 1]], #11 /32
        [-1, 1, nn.Upsample, [None, 2, "nearest"]], #12 /16
        [[-1, 6], 1, Concat, [1]], # cat backbone P4 /16       #13
        [-1, 3, C3, [512, False]], # 13 / 16                  #14
    
        [-1, 1, Conv, [256, 1, 1]], #15 /16 
        [-1, 1, nn.Upsample, [None, 2, "nearest"]], #16 /8
        [[-1, 4], 1, Concat, [1]], # cat backbone P3  /8    #17
        [-1, 3, C3, [256, False]], # 17 (P3/8-small) /8    #18
    
        [-1, 1, Conv, [256, 3, 2]], #19 /16
        [[-1, 6], 1, Concat, [1]], # cat head P4         #20
        [-1, 3, C3, [512, False]], # 20 (P4/16-medium)    #21
    
        [-1, 1, Conv, [512, 3, 2]], #22 /32
        [[-1, 8], 1, Concat, [1]], # cat head P5         #23
        [-1, 3, C3, [1024, False]], # 23 (P5/32-large)    #24
    
        [[18, 21, 24], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
      ]

在backbone最后一层添加了SELayer，这个类我已经在common.py中添加进来：

class SELayer(nn.Module):
        def __init__(self, c1, r=16):
            super(SELayer, self).__init__()
            self.avgpool = nn.AdaptiveAvgPool2d(1)
            self.l1 = nn.Linear(c1, c1//r, bias=False)
            self.relu = nn.ReLU(inplace=True)
            self.l2 = nn.Linear(c1//r, c1, bias=False)
            self.sig = nn.Sigmoid()
            
    
        def forward(self, x):
            b, c, _, _ = x.size()
            y = self.avgpool(x).view(b, c)
            y = self.l1(y)
            y = self.relu(y)
            y = self.l2(y)
            y = self.sig(y)
            y = y.view(b, c, 1, 1)
            return x * y.expand_as(x)

还需要在yolo.py中添加这个改动：

for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):
        m = eval(m) if isinstance(m, str) else m  # eval strings
        for j, a in enumerate(args):
            try:
                args[j] = eval(a) if isinstance(a, str) else a  # eval strings
            except:
                pass
        n = max(round(n * gd), 1) if n > 1 else n  # depth gain
        if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP,
                    DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, 
                    C3]:
            c1, c2 = ch[f], args[0]
            if c2 != no:  # if not output
                c2 = make_divisible(c2 * gw, 8)
    
            args = [c1, c2, *args[1:]]
            if m in [BottleneckCSP, C3]:
                args.insert(2, n)  # number of repeats
                n = 1
        elif m is nn.BatchNorm2d:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum([ch[x] for x in f])
        elif m is Detect:
            args.append([ch[x] for x in f])
            if isinstance(args[1], int):  # number of anchors
                args[1] = [list(range(args[1] * 2))] * len(f)
        elif m is Contract:
            c2 = ch[f] * args[0] ** 2
        elif m is Expand:
            c2 = ch[f] // args[0] ** 2
        elif m is SELayer: # 这里是修改的部分
            channel, re = args[0], args[1]
            channel = make_divisible(channel * gw, 8) if channel != no else channel 
            args = [channel, re]
        else:
            c2 = ch[f]