Pytorch实现CSPDarknet53

本是古典何须时尚 2023-01-13 10:58 206阅读 0赞

import torch
    import torch.nn.functional as F 
    import torch.nn as nn
    import math
    from collections import OrderedDict
    import numpy as np
    import torch.onnx
    
    # Mish激活函数: f(x) = x * tanh(log(1 + e^x))
    # pytorch中F.softplus(x)等同于torch.log(1 + torch.exp(x))
    class Mish(nn.Module):
        def __init__(self):
            super(Mish, self).__init__()
    
        def forward(self, x):
            # x * torch.tanh(torch.log(1 + torch.exp(x)))
            return x * torch.tanh(F.softplus(x))
    
    # 卷积块
    # conv + batchnormal + mish
    class BasicConv(nn.Module):
        def __init__(self, in_channels, out_channels, kernel_size, stride=1):
            super(BasicConv, self).__init__()
            
            # yolov4中只有1x1、3x3卷积, 1x1 padding=0, 3x3 padding=1, 卷积操作默认bias=False
            self.conv = nn.Conv2d(in_channels=in_channels, 
                                  out_channels=out_channels, 
                                  kernel_size=kernel_size, 
                                  stride=stride, 
                                  padding=kernel_size//2, 
                                  bias=False)
    
            self.bn = nn.BatchNorm2d(out_channels)
    
            self.activation = Mish()
            # self.activation = nn.ReLU()
        
        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):
            super(Resblock, self).__init__()
    
            if hidden_channels is None:
                hidden_channels = channels
            
            self.block = nn.Sequential(
                BasicConv(in_channels=channels, out_channels=hidden_channels, kernel_size=1),
                BasicConv(in_channels=hidden_channels, out_channels=channels, kernel_size=3),
            )
    
        def forward(self, x):
            return x + self.block(x)
    
    
    # CSPNet
    # 存在一个残差边，这个残差边绕过了很多残差结构
    class ResblockBody(nn.Module):
        def __init__(self, in_channels, out_channels, num_blocks, first):
            super(ResblockBody, self).__init__()
            
            self.downsample_conv = BasicConv(in_channels, out_channels, 3, stride=2)
    
            if first:
                self.split_conv0 = BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
                self.split_conv1 = BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
                self.blocks_conv = nn.Sequential(
                    Resblock(channels=out_channels, hidden_channels=out_channels//2),
                    BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
                )
    
                self.concat_conv = BasicConv(in_channels=out_channels * 2, out_channels=out_channels, kernel_size=1)
    
    
            else:
                self.split_conv0 = BasicConv(in_channels=out_channels, out_channels=out_channels//2, kernel_size=1)
                self.split_conv1 = BasicConv(in_channels=out_channels, out_channels=out_channels//2, kernel_size=1)
                
    
                self.blocks_conv = nn.Sequential(
                    *[Resblock(channels=out_channels//2) for _ in range(num_blocks)],
                    BasicConv(in_channels=out_channels//2, out_channels=out_channels//2, kernel_size=1)
                )
    
                self.concat_conv = BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
    
        def forward(self, x):                                                       
            x = self.downsample_conv(x)
            x0 = self.split_conv0(x)
            x1 = self.split_conv1(x)
            x1 = self.blocks_conv(x1)
    
            x = torch.cat([x1, x0], dim=1)
            x = self.concat_conv(x)
    
            return x
    
    # CSPDarknet
    class CSPDarknet53(nn.Module):
        def __init__(self, layers):
            super(CSPDarknet53, 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([
                ResblockBody(self.inplanes, self.feature_channels[0], layers[0], first=True),
                ResblockBody(self.feature_channels[0], self.feature_channels[1], layers[1], first=False),
                ResblockBody(self.feature_channels[1], self.feature_channels[2], layers[2], first=False),
                ResblockBody(self.feature_channels[2], self.feature_channels[3], layers[3], first=False),
                ResblockBody(self.feature_channels[3], self.feature_channels[4], layers[4], first=False)
            ])
    
            # 进行权值初始化
            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)
            out4 = self.stages[3](out3)
            out5 = self.stages[4](out4)
    
            return out3, out4, out5
    
    
    def load_model(model, pth):
        # 获取可用设备
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
        # 获取网络参数字典(参数键值对)
        model_dict = model.state_dict()
    
        # 获取预训练模型参数字典(参数键值对)
        pretrained_dict = torch.load(pth, map_location=device)
        for k, v in pretrained_dict.items():
            # print(k)
            pass
    
        # 网络参数与预训练模型匹配的参数保存至字典matched_dict中
        matched_dict = {}
        for k, v in model_dict.items():
            # print(k)
            if k.find('backbone') == -1:
                # 预训练模型的键存在backbon.字段
                key = 'backbone.' + k
                # 比较预训练模型与构建网络参数维度是否一致
                if np.shape(pretrained_dict[key]) == np.shape(v):
                    # 将预训练模型参数赋值给构建的网络
                    matched_dict[k] = pretrained_dict[key]
    
        for key in matched_dict:
            print('pretrained items:', key)
        
        print('%d layers matched, %d layers miss'%(len(matched_dict.keys()), len(model_dict) - len(matched_dict.keys())))
        model_dict.update(matched_dict)
        model.load_state_dict(model_dict)
        print('Finished !')
        return model
    
    if __name__ == "__main__":
        net = CSPDarknet53([1, 2, 8, 8, 4])
        model = load_model(net, './weights/CSPDarknet53.pth')
        model.eval()
        data = torch.randn(1, 3, 608, 608)
        torch.onnx.export(model, data, "CSPDarknet53.onnx", verbose=True, input_names = ['input'], output_names = ['output'])

生成简化的onnx模型：

python -m onnxsim CSPDarknet53.onnx CSPDarknet53Sim.onnx

模型文件下载：[https://download.csdn.net/download/weicao1990/16745420?spm=1001.2014.3001.5501][https_download.csdn.net_download_weicao1990_16745420_spm_1001.2014.3001.5501]

参考：[https://blog.csdn.net/weixin\_41560402/article/details/106119774][https_blog.csdn.net_weixin_41560402_article_details_106119774]

[https_download.csdn.net_download_weicao1990_16745420_spm_1001.2014.3001.5501]: https://download.csdn.net/download/weicao1990/16745420?spm=1001.2014.3001.5501
[https_blog.csdn.net_weixin_41560402_article_details_106119774]: https://blog.csdn.net/weixin_41560402/article/details/106119774