专栏名称: OpenCV学堂

一个致力于计算机视觉OpenCV原创技术传播的公众号！OpenCV计算机视觉与tensorflow深度学习相关算法原创文章分享、函数使用技巧、源码分析与讨论、，计算机视觉前沿技术介绍，技术专家经验分享，人才交流，学习交流。

YOLOv5魔改替换感受野注意力(RFAConv)模块从训练到推理

OpenCV学堂 · 公众号 · · 2024-05-10 22:28

正文

点击上方蓝字关注我们

微信公众号： OpenCV学堂

关注获取更多计算机视觉与深度学习知识

YOLOv5结构与源码修改

以YOLOv5n.yaml 文件结构为例，替换其中的下采样Conv为RFAConv，模型结构修改为：

然后在common.py文件添加RFAConv的代码实现如下：

class RFAConv(nn.Module):
    def __init__(self,in_channel, out_channel, kernel_size=1, stride=1):
        super().__init__()
        self.kernel_size = kernel_size
        self.get_weight = nn.Sequential(nn.AvgPool2d(kernel_size=kernel_size, padding=kernel_size // 2, stride=stride),
                                        nn.Conv2d(in_channel, in_channel * (kernel_size ** 2), kernel_size=1, groups=in_channel,bias=False))
        self.generate_feature = nn.Sequential(
            nn.Conv2d(in_channel, in_channel * (kernel_size ** 2), kernel_size=kernel_size,padding=kernel_size//2,stride=stride, groups=in_channel, bias=False),
            nn.BatchNorm2d(in_channel * (kernel_size ** 2)),
            nn.ReLU())

        self.conv = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=kernel_size, stride=kernel_size),
                                  nn.BatchNorm2d(out_channel),
                                  nn.ReLU())

    def forward(self,x):
        b,c = x.shape[0:2]
        weight =  self.get_weight(x)
        h,w = weight.shape[2:]
        weighted = weight.view(b, c, self.kernel_size ** 2, h, w).softmax(2)  # b c*kernel**2,h,w ->  b c k**2 h w 
        feature = self.generate_feature(x).view(b, c, self.kernel_size ** 2, h, w)  #b c*kernel**2,h,w ->  b c k**2 h w   获得感受野空间特征
        weighted_data = feature * weighted
        conv_data = rearrange(weighted_data, 'b c (n1 n2) h w -> b c (h n1) (w n2)', n1=self.kernel_size, # b c k**2 h w ->  b c h*k w*k
                              n2=self.kernel_size)
        return self.conv(conv_data)

最后在yolo.py中完成RFAConv模块注册。

模型训练与ONNX导出

基于OID数据集中的大象与骆驼数据集，用下面的命令行开启YOLOv5 + RFAConv修改后的模型训练。注意这里weights参数必须设置为空，意思是从指定的yaml文件开始从零训练模型，命令行如下：

python train.py –weights ‘’ –cfg ./model/yolov5n-rfa.yaml –data ./camel_elephant_training/dataset.yaml –epochs 100 –batch-size 4

运行完成以后执行，看到的PR曲线

导出ONNX格式模型，执行的命令行如下：

python export.py –weights elemphant_rfa_best.pt –include onnx

可以看出导出以后的模型非常小，文件大小只有5MB，真很不错。YOLOv5n魔改也很靠谱。

模型推理演示

基于QT的推理测试代码如下：

YOLOv5魔改替换感受野注意力(RFAConv)模块从训练到推理

正文

请到「今天看啥」查看全文