YOLOV5目标检测换多种主干的方法（含shufflenet，ghost net，mobilenet，resnet等）
一，YOLOv5主干换为ghostnet
其实换主干是一个较为简单的操作方法，但是很多新手可能不太会，或者其中可能会出现一些问题，那我对这些问题归纳出三步走换主干的方法，以便和大家分享交流。我们这里以yolov5s为例。
1，找到想要替换的新的主干网络的基本结构组成，并将其复制到models/common.py文件中。与yolov5拥有C3，Conv这些基本组成结构相似，ghostnet也有类似的基本结构，它们分别是Ghostconv，Ghostbottleneck和C3Ghost，实际上他们对应的就是主干的不同级别的单位组成。将它们放置在models/common.py中的150行以后，找个位置插入即可。
class GhostConv(nn.Module):
def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
super(GhostConv, self).__init__()
c_ = c2 // 2
self.cv1 = Conv(c1, c_, k, s, None, g, act)
self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
def forward(self, x):
y = self.cv1(x)
return torch.cat([y, self.cv2(y)], 1)
class GhostBottleneck(nn.Module):
def __init__(self, c1, c2, k=3, s=1):
super().__init__()
c_ = c2 // 2
self.conv = nn.Sequential(
GhostConv(c1, c_, 1, 1),
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(),
GhostConv(c_, c2, 1, 1, act=False))
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1,
act=False)) if s == 2 else nn.Identity()
def forward(self, x):
return self.conv(x) + self.shortcut(x)
class C3Ghost(C3):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e)
self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
2，在结构搭建脚本models/yolo.py添加上我们新增加的模块名称，以便搭建网络时可以直接调用。不同的版本这一块的内容可能不完全相同，我们要做的就是把Ghostconv，Ghostbottleneck和C3Ghost补全进去即可。
3，接下来就是核心步骤了，大家可以直接在models/yolov5s.yaml里面进行修改，或者重命名一个文件，但重命名的话记得在train.py里面改一下路径，这里就简单的直接进行修改。因为我们只修改主干，因此只需要改动backbone这里的内容即可。实际上来讲，这里可操作的空间很大，大家只要保证P3，P4，P5层的下采样倍数和通道数是没问题的就行，其他中间具体的结构其实改了也不会报错，所以有兴趣的小伙伴不如炼丹试试，下面我分享自己进行改进的内容：
之后直接去运行train.py即可。
二，yolov5主干换为shufflenet
1，这里依旧是三步走的步骤，对common文件进行模块内容的添加。
def channel_shuffle(x, groups):
batchsize, num_channels, height, width = x.data.size()
channels_per_group = num_channels // groups
x = x.view(batchsize, groups, channels_per_group, height, width)
x = torch.transpose(x, 1, 2).contiguous()
x = x.view(batchsize, -1, height, width)
return x
class stem(nn.Module):
def __init__(self, c1, c2):
super(stem, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(c1, c2, kernel_size=3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(c2),
nn.ReLU(inplace=True),
)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
def forward(self, x):
return self.maxpool(self.conv(x))
class Shuffle_Block(nn.Module):
def __init__(self, ch_in, ch_out, stride):
super(Shuffle_Block, self).__init__()
if not (1 <= stride <= 2):
raise ValueError('illegal stride value')
self.stride = stride
branch_features = ch_out // 2
assert (self.stride != 1) or (ch_in == branch_features << 1)
if self.stride > 1:
self.branch1 = nn.Sequential(
self.depthwise_conv(ch_in, ch_in, kernel_size=3, stride=self.stride, padding=1),
nn.BatchNorm2d(ch_in),
nn.Conv2d(ch_in, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(branch_features),
nn.ReLU(inplace=True),
)
self.branch2 = nn.Sequential(
nn.Conv2d(ch_in if (self.stride > 1) else branch_features,
branch_features, kernel_size=1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(branch_features),
nn.ReLU(inplace=True),
self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1),
nn.BatchNorm2d(branch_features),
nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(branch_features),
nn.ReLU(inplace=True),
)
@staticmethod
def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):
return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)
def forward(self, x):
if self.stride == 1:
x1, x2 = x.chunk(2, dim=1)
out = torch.cat((x1, self.branch2(x2)), dim=1)
else:
out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
out = channel_shuffle(out, 2)
return out
2，对models/yolo.py进行添加这几个模块名称。
3，修改models/yolov5s.yaml文件内容
这就完成了。
三，yolov5主干改为resnet
1，仍然是添加新的模块进入models/common.py,你们用的时候尽量直接复制，尽量不要修改。
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channel)
self.relu = nn.ReLU()
self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channel)
self.downsample = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
kernel_size=3, stride=2, padding=1, bias=False)
def forward(self, x):
identity = x
if self.downsample is not None:
identity = self.downsample(x)
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out += identity
out = self.relu(out)
return out
class Bottleneck1(nn.Module):
"""
注意：原论文中，在虚线残差结构的主分支上，第一个1x1卷积层的步距是2，第二个3x3卷积层步距是1。
但在pytorch官方实现过程中是第一个1x1卷积层的步距是1，第二个3x3卷积层步距是2，
这么做的好处是能够在top1上提升大概0.5%的准确率。
可参考Resnet v1.5 https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch
"""
expansion = 1
def __init__(self, in_channel, out_channel, stride=1, downsample=None,
groups=1, width_per_group=64):
super(Bottleneck1, self).__init__()
width = int(out_channel * (width_per_group / 64.)) * groups
self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=width,
kernel_size=1, stride=1, bias=False) # squeeze channels
self.bn1 = nn.BatchNorm2d(width)
# -----------------------------------------
self.conv2 = nn.Conv2d(in_channels=width, out_channels=width, groups=groups,
kernel_size=3, stride=stride, bias=False, padding=1)
self.bn2 = nn.BatchNorm2d(width)
# -----------------------------------------
self.conv3 = nn.Conv2d(in_channels=width, out_channels=out_channel*self.expansion,
kernel_size=1, stride=1, bias=False) # unsqueeze channels
self.bn3 = nn.BatchNorm2d(out_channel*self.expansion)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
def forward(self, x):
identity = x
if self.downsample is not None:
identity = self.downsample(x)
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
out += identity
out = self.relu(out)
return out2，添加模块名称，你们注意这里是Bottleneck1，不要忘了后面有1.
3，修改配置文件，我并没有严格按照resnet的层数堆，如果想要很标准的话，可以自己查一下每个模块堆叠多少次。
