0
点赞
收藏
分享

微信扫一扫

YOLOv5:修改backbone为SPD-Conv

YOLOv5:修改backbone为SPD-Conv

在这里插入图片描述

前言

前提条件

相关介绍

SPD-Conv

在这里插入图片描述
以下是使用Python实现SPD-Conv的简单例子,目的是方便大家理解SPD-Conv的操作。

import cv2
import torch
from torch import nn

############## SPD-Conv ##############
class space_to_depth(nn.Module):
    # Changing the dimension of the Tensor
    def __init__(self, dimension=1):
        super().__init__()
        self.d = dimension

    def forward(self, x):
         return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
############## SPD-Conv ##############


if __name__=="__main__":

    img_tensor = torch.Tensor([[[[ 1, 2, 3, 4],
                                 [ 5, 6, 7, 8],
                                 [ 9, 10, 11, 12],
                                 [ 13, 14, 15, 16]],

                                 [[ 1, 2, 3, 4],
                                 [ 5, 6, 7, 8],
                                 [ 9, 10, 11, 12],
                                 [ 13, 14, 15, 16]],

                                 [[ 1, 2, 3, 4],
                                 [ 5, 6, 7, 8],
                                 [ 9, 10, 11, 12],
                                 [ 13, 14, 15, 16]]]])
    # print('img_tensor:',img_tensor)
    print('img_tensor.shape:',img_tensor.shape)
    spd = space_to_depth()
    res = spd.forward(img_tensor)
    # print('res:',res)
    print('res.shape:',res.shape)
img_tensor.shape: torch.Size([1, 3, 4, 4])
res: tensor([[[[ 1.,  3.],
       		[ 9., 11.]],
	
	         [[ 1.,  3.],
	          [ 9., 11.]],
	
	         [[ 1.,  3.],
	          [ 9., 11.]],
	
	         [[ 5.,  7.],
	          [13., 15.]],
	
	         [[ 5.,  7.],
	          [13., 15.]],
	
	         [[ 5.,  7.],
	          [13., 15.]],
	
	         [[ 2.,  4.],
	          [10., 12.]],
	
	         [[ 2.,  4.],
	          [10., 12.]],
	
	         [[ 2.,  4.],
	          [10., 12.]],
	
	         [[ 6.,  8.],
	          [14., 16.]],
	
	         [[ 6.,  8.],
	          [14., 16.]],
	
	         [[ 6.,  8.],
	          [14., 16.]]]])
res.shape: torch.Size([1, 12, 2, 2])

YOLOv5修改backbone为SPD-Conv

在这里插入图片描述

修改common.py

将以下代码,添加进common.py。

############## SPD-Conv ##############
class space_to_depth(nn.Module):
    # Changing the dimension of the Tensor
    def __init__(self, dimension=1):
        super().__init__()
        self.d = dimension

    def forward(self, x):
         return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
############## SPD-Conv ##############

修改yolo.py

        elif m is space_to_depth:
            c2 = 4 * ch[f]

在这里插入图片描述

修改yolov5.yaml配置

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
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 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],     # 0-P1/2
   [-1, 1, Conv, [128, 3, 1]],  # 1
   [-1,1,space_to_depth,[1]],   # 2 -P2/4
   [-1, 3, C3, [128]],          # 3
   [-1, 1, Conv, [256, 3, 1]],  # 4
   [-1,1,space_to_depth,[1]],   # 5 -P3/8
   [-1, 6, C3, [256]],          # 6
   [-1, 1, Conv, [512, 3, 1]],  # 7-P4/16
   [-1,1,space_to_depth,[1]],   # 8 -P4/16
   [-1, 9, C3, [512]],          # 9
   [-1, 1, Conv, [1024, 3, 1]], # 10-P5/32
   [-1,1,space_to_depth,[1]],   # 11 -P5/32
   [-1, 3, C3, [1024]],         # 12
   [-1, 1, SPPF, [1024, 5]],    # 13
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],                    # 14
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],    # 15
   [[-1, 9], 1, Concat, [1]],                     # 16 cat backbone P4
   [-1, 3, C3, [512, False]],                     # 17

   [-1, 1, Conv, [256, 1, 1]],                    # 18
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],    # 19
   [[-1, 6], 1, Concat, [1]],                     # 20 cat backbone P3
   [-1, 3, C3, [256, False]],                     # 21 (P3/8-small)

   [-1, 1, Conv, [256, 3, 1]],                    # 22
   [-1,1,space_to_depth,[1]],                     # 23 -P2/4
   [[-1, 18], 1, Concat, [1]],                    # 24 cat head P4
   [-1, 3, C3, [512, False]],                     # 25 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 1]],                    # 26
   [-1,1,space_to_depth,[1]],                     # 27 -P2/4
   [[-1, 14], 1, Concat, [1]],                    # 28 cat head P5
   [-1, 3, C3, [1024, False]],                    # 29 (P5/32-large)

   [[21, 25, 29], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

在这里插入图片描述

参考

[1] Raja Sunkara, Tie Luo. No More Strided Convolutions or Pooling: A New CNN Building Block for Low-Resolution Images and Small Objects. 2022
[2] https://github.com/LabSAINT/SPD-Conv
[3] https://github.com/ultralytics/yolov5.git

举报

相关推荐

0 条评论