How can I add new backbone network for SSD and DSSD
sd59202 opened this issue · 0 comments
I try to add darknet53 for backbone and modify the network but not work.
Please help to check how can I add this SSD & ZQPei/DSSD network.
``
import torch
from torch import nn
import math
from dssd.utils.model_zoo import load_state_dict_from_url
from collections import OrderedDict
from torchsummary import summary
all = ['darknet53']
model_urls = {
'darknet53':"https://github.com/Jintao-Huang/Darknet53_PyTorch/releases/download/1.0/darknet53-26b80406.pth"
}
def load_pretrained(model, state_dict):
model.load_state_dict(state_dict, strict=False)
def _darknet(arch, layers, pretrained, progress, **kwargs):
model = DarkNet(layers)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch])
load_pretrained(model, state_dict)
return model
class BasicBlock(nn.Module):
def init(self, inplanes, planes):
super(BasicBlock, self).init()
self.conv1 = nn.Conv2d(inplanes, planes[0], kernel_size=1,
stride=1, padding=0, bias=False)
self.bn1 = nn.BatchNorm2d(planes[0])
self.relu1 = nn.LeakyReLU(0.1)
self.conv2 = nn.Conv2d(planes[0], planes[1], kernel_size=3,
stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes[1])
self.relu2 = nn.LeakyReLU(0.1)
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu1(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu2(out)
out += residual
return out
class DarkNet(nn.Module):
def init(self, layers):
super(DarkNet, self).init()
self.inplanes = 32
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(self.inplanes)
self.relu1 = nn.LeakyReLU(0.1)
self.layer1 = self._make_layer([32, 64], layers[0])
self.layer2 = self._make_layer([64, 128], layers[1])
self.layer3 = self._make_layer([128, 256], layers[2])
self.layer4 = self._make_layer([256, 512], layers[3])
self.layer5 = self._make_layer([512, 1024], layers[4])
self.layer6 = self._make_extra_layer([1024, 1024])
self.layer7 = self._make_extra_layer([1024, 1024])
self.layer8 = self._make_extra_layer([1024, 1024])
self.layer9 = self._make_extra_layer([1024, 1024])
self.layers_out_filters = [64, 128, 256, 512, 1024]
self.init_params()
# initialize the parameters in convolution and batch normalization layers
def init_params(self):
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 _make_layer(self, planes, blocks):
layers = []
# downsample
layers.append(("ds_conv", nn.Conv2d(self.inplanes, planes[1], kernel_size=3,
stride=2, padding=1, bias=False)))
layers.append(("ds_bn", nn.BatchNorm2d(planes[1])))
layers.append(("ds_relu", nn.LeakyReLU(0.1)))
# blocks
self.inplanes = planes[1]
for i in range(0, blocks):
layers.append(("residual_{}".format(i), BasicBlock(self.inplanes, planes)))
return nn.Sequential(OrderedDict(layers))
def _make_extra_layer(self, planes):
layers = []
# downsample
layers.append(("ds_conv", nn.Conv2d(self.inplanes, planes[1], kernel_size=3,
stride=2, padding=1, bias=False)))
layers.append(("ds_bn", nn.BatchNorm2d(planes[1])))
layers.append(("ds_relu", nn.LeakyReLU(0.1)))
# blocks
self.inplanes = planes[1]
for i in range(0, 1):
layers.append(("residual_{}".format(i), BasicBlock(self.inplanes, planes)))
return nn.Sequential(OrderedDict(layers))
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
out1 = self.layer1(x)
out2 = self.layer2(out1)
out3 = self.layer3(out2)
out4 = self.layer4(out3)
out5 = self.layer5(out4)
out6 = self.layer6(out5)
out7 = self.layer7(out6)
out8 = self.layer8(out7)
out9 = self.layer9(out8)
return out3, out4, out5, out6, out7, out8, out9
def darknet21(pretrained, **kwargs):
"""Constructs a darknet-21 model.
"""
model = DarkNet([1, 1, 2, 2, 1])
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
def darknet53(pretrained=False, progress=True, **kwargs):
"""Constructs a darknet-53 model.
"""
#model = DarkNet([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
return _darknet('darknet53', [1, 2, 8, 8, 4], pretrained, progress, **kwargs)
if name == 'main':
darknet = DarkNet([1,2,8,8,4]).cuda()
summary(darknet, (3,320,320))
``