pytorch实现yolov4CSPDarknet53主干网络


```python
import torch
import torch.nn.functional as F
import torch.nn as nn
import math
from collections import OrderedDict

#-------------------------------------------------#
#   MISH激活函数
#-------------------------------------------------#
class Mish(nn.Module):
    def __init__(self):
        super(Mish, self).__init__()

    def forward(self, x):
        return x * torch.tanh(F.softplus(x))

#-------------------------------------------------#
#   卷积块
#   CONV+BATCHNORM+MISH
#-------------------------------------------------#
class BasicConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
        super(BasicConv, self).__init__()

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.activation = Mish()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.activation(x)
        return x

#---------------------------------------------------#
#   CSPdarknet的结构块的组成部分
#   内部堆叠的残差块
#---------------------------------------------------#
class Resblock(nn.Module):
    def __init__(self, channels, hidden_channels=None, residual_activation=nn.Identity()):
        super(Resblock, self).__init__()

        if hidden_channels is None:
            hidden_channels = channels

        self.block = nn.Sequential(
            BasicConv(channels, hidden_channels, 1),
            BasicConv(hidden_channels, channels, 3)
        )

    def forward(self, x):
        return x + self.block(x)

#---------------------------------------------------#
#   CSPdarknet的结构块
#   存在一个大残差边
#   这个大残差边绕过了很多的残差结构
#---------------------------------------------------#
class Resblock_body(nn.Module):
    def __init__(self, in_channels, out_channels, num_blocks, first):
        super(Resblock_body, self).__init__()

        self.downsample_conv = BasicConv(in_channels, out_channels, 3, stride=2)

        if first:
            self.split_conv0 = BasicConv(out_channels, out_channels, 1) #csp net part1
            self.split_conv1 = BasicConv(out_channels, out_channels, 1)  #csp net part2
            self.blocks_conv = nn.Sequential(#csp net part2
                Resblock(channels=out_channels, hidden_channels=out_channels//2),
                BasicConv(out_channels, out_channels, 1)
            )
            self.concat_conv = BasicConv(out_channels*2, out_channels, 1) #csp net 拼接
        else:
            self.split_conv0 = BasicConv(out_channels, out_channels//2, 1)
            self.split_conv1 = BasicConv(out_channels, out_channels//2, 1)

            self.blocks_conv = nn.Sequential(
                *[Resblock(out_channels//2) for _ in range(num_blocks)], #多个resblock循环叠加
                BasicConv(out_channels//2, out_channels//2, 1)
            )
            self.concat_conv = BasicConv(out_channels, out_channels, 1)

    def forward(self, x):
        x = self.downsample_conv(x)
        # csp net part1
        x0 = self.split_conv0(x)
        # csp net part2
        x1 = self.split_conv1(x)
        x1 = self.blocks_conv(x1)
        #csp 通道拼接
        x = torch.cat([x1, x0], dim=1)
        #卷积
        x = self.concat_conv(x)

        return x

class CSPDarkNet(nn.Module):
    def __init__(self, layers):#layers :[1,2,8,8,4] 每个Resblock_body的重复个数
        super(CSPDarkNet, self).__init__()
        self.inplanes = 32
        self.conv1 = BasicConv(3, self.inplanes, kernel_size=3, stride=1)
        self.feature_channels = [64, 128, 256, 512, 1024]

        self.stages = nn.ModuleList([
            Resblock_body(self.inplanes, self.feature_channels[0], layers[0], first=True),
            Resblock_body(self.feature_channels[0], self.feature_channels[1], layers[1], first=False),#输入64维,输出128维,层数2层,
            Resblock_body(self.feature_channels[1], self.feature_channels[2], layers[2], first=False),
            Resblock_body(self.feature_channels[2], self.feature_channels[3], layers[3], first=False),
            Resblock_body(self.feature_channels[3], self.feature_channels[4], layers[4], first=False)
        ])

        self.num_features = 1
        # 进行权值初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


    def forward(self, x):
        x = self.conv1(x)

        x = self.stages[0](x)
        x = self.stages[1](x)
        out3 = self.stages[2](x) #8倍采样
        out4 = self.stages[3](out3) #16背下采样
        out5 = self.stages[4](out4) #32倍下采样

        return out3, out4, out5

def darknet53(pretrained, **kwargs):
    model = CSPDarkNet([1, 2, 8, 8, 4])
    if pretrained:
        if isinstance(pretrained, str):
            model.load_state_dict(torch.load(pretrained))
        else:
            raise Exception("darknet request a pretrained path. got [{}]".format(pretrained))
    return model


CSPDarknet53主干网络中添加注意力机制的方法如下: 1. 首先,在CSPDarknet53主干网络的每个卷积层之后添加一个注意力模块。注意力模块是一种能够自适应地学习输入特征图中不同区域的重要性的模块。 2. 注意力模块通常由一个全局平均池化层、一个全连接层和一个sigmoid激活函数组成。全局平均池化层用于对输入特征图进行降维,全连接层用于学习特征图中每个位置的重要性,sigmoid激活函数用于将注意力权重限制在0到1之间。 3. 在实现过程中,可以使用PyTorch框架中的torch.nn模块来实现注意力模块。具体实现步骤如下: 1)定义一个Attention模块,该模块包括全局平均池化层、全连接层和sigmoid激活函数。 ```python class Attention(nn.Module): def __init__(self, in_channels): super(Attention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(in_channels, in_channels // 16, bias=False), nn.ReLU(inplace=True), nn.Linear(in_channels // 16, in_channels, bias=False), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) y = self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x) ``` 2)在CSPDarknet53主干网络的每个卷积层之后添加一个Attention模块。 ```python class CSPDarknet53(nn.Module): def __init__(self): super(CSPDarknet53, self).__init__() # 省略其他代码 self.conv6 = nn.Conv2d(512, 1024, 3, stride=2, padding=1) self.bn6 = nn.BatchNorm2d(1024) self.attention6 = Attention(1024) # 省略其他代码 def forward(self, x): # 省略其他代码 x = self.conv6(x) x = self.bn6(x) x = self.attention6(x) # 省略其他代码 return x ``` 4. 最后,在训练过程中,使用交叉熵损失函数来训练模型,并使用随机梯度下降(SGD)或Adam优化器来更新模型参数。 这样,就可以在CSPDarknet53主干网络中添加注意力机制了。注意力机制可以帮助模型更好地学习输入特征图中不同区域的重要性,从而提高模型的性能。
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值