flyfish
版本:YOLOv5:v5
具体已经借鉴的自定义网络结构包括
MobileNetV3 Large
MobileNetV3 Small
ShuffleNetV2
ShuffleNetV2-Focus
ShuffleNetV2-stem(Pelee的stem模块)
python train.py --batch 64 --epochs 300 --data data/coco128.yaml --cfg models/yolov5-mobilenetv3small.yaml
yolov5-mobilenetv3small.yaml 可以更换为如下配置
├── yolov5-mobilenetv3large.yaml
├── yolov5-mobilenetv3small.yaml
├── yolov5-shufflenetv2-focus.yaml
├── yolov5-shufflenetv2-stem.yaml
├── yolov5-shufflenetv2.yaml
一种方式是可以将原结构整理成
backbones
necks
heads
detectors
losses
utils
例如backbones里面有shufflenetv2,mobilenetv3.py等
另一种方式
采用YOLOv5的配置方式
添加模块,根据配置文件生成网络结构。
这里使用的是第二种方式
本文以将ShuffleNetV2的InvertedResidual模块加入到YOLOv5的Backbone中为例说明如何自定义网络结构,加入自定义模块
第一步:加入模块代码
在common.py的顶部加入导入
from torch import Tensor
from typing import Callable, Any, List
将InvertedResidual类和InvertedResidual类需要的channel_shuffle函数都加入到common.py的底部
def channel_shuffle(x: Tensor, groups: int) -> Tensor:
batchsize, num_channels, height, width = x.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 InvertedResidual(nn.Module):
def __init__(
self,
inp: int,
oup: int,
stride: int
) -> None:
super(InvertedResidual, self).__init__()
if not (1 stride 3):
raise ValueError('illegal stride value')
self.stride = stride
branch_features = oup // 2
assert (self.stride != 1) or (inp == branch_features << 1)
if self.stride > 1:
self.branch1 = nn.Sequential(
self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1),
nn.BatchNorm2d(inp),
nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(branch_features),
nn.ReLU(inplace=True),
)
else:
self.branch1 = nn.Sequential()
self.branch2 = nn.Sequential(
nn.Conv2d(inp 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: int,
o: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
bias: bool = False
) -> nn.Conv2d:
return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)
def forward(self, x: Tensor) -> Tensor:
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
第二步 更改解析模块,告诉YOLOv5,我们加入了InvertedResidual模块
目录在models/yolo.py的parse_model函数
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
C3, C3TR,InvertedResidual]:
第三步配置
在目录models下新建yolov5-shufflenetv2-focus.yaml文件,配置如下
nc: 80
depth_multiple: 1.0
width_multiple: 0.5
anchors:
- [4,5, 8,10, 13,16]
- [23,29, 43,55, 73,105]
- [146,217, 231,300, 335,433]
backbone:
[[-1, 1, Focus, [64, 3]],
[-1, 1, InvertedResidual, [128, 2]],
[-1, 3, InvertedResidual, [128, 1]],
[-1, 1, InvertedResidual, [256, 2]],
[-1, 7, InvertedResidual, [256, 1]],
[-1, 1, InvertedResidual, [512, 2]],
[-1, 3, InvertedResidual, [512, 1]],
]
head:
[[-1, 1, Conv, [128, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]],
[-1, 1, C3, [128, False]],
[-1, 1, Conv, [128, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 2], 1, Concat, [1]],
[-1, 1, C3, [128, False]],
[-1, 1, Conv, [128, 3, 2]],
[[-1, 11], 1, Concat, [1]],
[-1, 1, C3, [128, False]],
[-1, 1, Conv, [128, 3, 2]],
[[-1, 7], 1, Concat, [1]],
[-1, 1, C3, [128, False]],
[[14, 17, 20], 1, Detect, [nc, anchors]],
]
配置的参数说明
在编写配置文件时,不需要定义输入通道,只需要定义输出通道和其他参数。
输入数据最初是3通道,定义输出通道,该通道也是其他层的输入,
例如
[[-1, 1, conv_bn_hswish, [16, 2]],
[-1, 1, MobileNetV3_InvertedResidual, [16, 16, 3, 2, 1, 0]],
[-1, 1, MobileNetV3_InvertedResidual, [24, 72, 3, 2, 0, 0]],
代码会自动扩展
models.common.conv_bn_hswish [3, 16, 2]
models.common.MobileNetV3_InvertedResidual [16, 16, 16, 3, 2, 1, 0]
models.common.MobileNetV3_InvertedResidual [24, 24, 88, 3, 1, 0, 0]
conv_bn_hswish的参数
self, c1, c2, stride
含义如下
输入通道是c1=3
输出通道是c2=16
stride=2
MobileNetV3_InvertedResidual的参数
self, inp, oup, hidden_dim, kernel_size, stride, use_se, use_hs
输入通道是inp=16,
输出通道是oup=16,
hidden_dim=16,
kernel_size=3,
stride=2,
use_se=1,
use_hs=0
use_se表示是否使用SELayer
use_hs表示使用h_swish还是ReLU
Original: https://blog.csdn.net/flyfish1986/article/details/117303291
Author: 西西弗Sisyphus
Title: 目标检测 YOLOv5 自定义网络结构
原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/686954/
转载文章受原作者版权保护。转载请注明原作者出处!