DarkNet53Pytorch实现和.pth的预训练权重下载

DarkNet53是Yolov3的主干网，当我们想拿来做分割或者分类的时候需要将其单独编写出来，并加载预训练的权重。我在网上找了挺久，不知道为什么权重文件都是.weights或者.conv结尾的，这样的文件貌似pytorch无法直接加载，所以本文给大家分享一下它的预训练权重，大家有需要的可以来下载：
链接：https://pan.baidu.com/s/1n6PRMSQGtdWzQqkyYTDc5w
提取码：3h03
关于代码，我是在别人的基础上进行了一些修改，使得最后两个block的步长为1，最后两个block卷积的空洞率分别为2和4,这样能扩大感受野。如果大家想要改回去，只需要将最后两个block的参数改乘strid为2和dilated为1即可。

"pytorch darcknet53"
"需要对其进行修改，用来做语义分割的主干网络"
import torch
import torch.nn as nn
import numpy as np
def Conv3x3BNReLU(in_channels,out_channels,stride=1,dilated=1):
    return nn.Sequential(
        nn.Conv2d(in_channels=in_channels,out_channels=out_channels,
                  kernel_size=3,stride=stride,padding=dilated,dilation=dilated,bias=False),
        nn.BatchNorm2d(out_channels),
        nn.ReLU6(inplace=True)
    )

def Conv1x1BNReLU(in_channels,out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels=in_channels,out_channels=out_channels,
                  kernel_size=1,stride=1,padding=0,bias=False),
        nn.BatchNorm2d(out_channels),
        nn.ReLU6(inplace=True)
    )

class Residual(nn.Module):
    def __init__(self, nchannels,dilated=1):
        super(Residual, self).__init__()
        mid_channels = nchannels // 2
        self.conv1x1 = Conv1x1BNReLU(in_channels=nchannels, out_channels=mid_channels)
        self.conv3x3 = Conv3x3BNReLU(in_channels=mid_channels, out_channels=nchannels,dilated=dilated)

    def forward(self, x):
        out = self.conv3x3(self.conv1x1(x))
        return out + x

class Darknet53(nn.Module):
    def __init__(self, num_classes=1000):
        super(Darknet53, self).__init__()
        self.first_conv = Conv3x3BNReLU(in_channels=3, out_channels=32)
        self.block1 = self._make_layers(in_channels=32,out_channels=64, block_num=1,stride=2)
        self.block2 = self._make_layers(in_channels=64,out_channels=128, block_num=2,stride=2)
        self.block3 = self._make_layers(in_channels=128,out_channels=256, block_num=8,stride=2)
        self.block4 = self._make_layers(in_channels=256,out_channels=512, block_num=8,stride=1,dilated=2)
        self.block5 = self._make_layers(in_channels=512,out_channels=1024, block_num=4,stride=1,dilated=4)

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.linear = nn.Linear(in_features=1024,out_features=num_classes)
        self.softmax = nn.Softmax(dim=1)

    def _make_layers(self, in_channels,out_channels, block_num,dilated=1,stride=1):
        _layers = []
        _layers.append(Conv3x3BNReLU(in_channels=in_channels, out_channels=out_channels, stride=stride,dilated=dilated))
        for _ in range(block_num):
            _layers.append(Residual(nchannels=out_channels,dilated=dilated))
        return nn.Sequential(*_layers)

    def forward(self, x):
        x = self.first_conv(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x) # 1 256 32 32
        x = self.block4(x)
        x = self.block5(x)

        x = self.avg_pool(x)
        x = x.view(x.size(0),-1)
        x = self.linear(x)
        out = self.softmax(x)
        return out

关于加载权重，你直接加载我提供的权重肯定是加载不上去的，因为存权重的时候的key不一样，当你需要加载一个需要预训练的网络时候，可以调用我下面这段代码，加载模型权重：

path = r'D:\1Apython\Pycharm_pojie\Semantic_segmentation\pytorch_segmentation-master\pretrained\darknet53-a628ea1b.pth'
def darknet53(pretrained=True,root=path):
    model = Darknet53()
    model_dict = model.state_dict()  # 拿到我们自己定义网络的权重参数
    if pretrained:
        ckpt = torch.load(root)
        new_state_dict = {}
        for (k,v),k2 in zip(ckpt['state_dict'].items(),model_dict.keys()):
            new_state_dict[k2] = v
        try:
            model.load_state_dict(new_state_dict,strict=False)
            print('darknet53 成功加载模型')
        except KeyError as e:
            s = "初始化权值字典加载错误"
            raise KeyError(s) from e
        del ckpt,new_state_dict
        return model
    return model
if __name__ == '__main__':
    model = darknet53()
    print(model)
    input = torch.randn(1,3,256,256)
    out = model(input)
    print(out.shape)

其中这个path为你存放权重的路径。