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为你存放权重的路径。