lufficc / SSD

High quality, fast, modular reference implementation of SSD in PyTorch
MIT License
1.52k stars 384 forks source link

How can I add new backbone network for SSD and DSSD #193

Closed sd59202 closed 2 years ago

sd59202 commented 3 years ago

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