lufficc/SSD

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))
``