retinanet.py

import torch.nn as nn
import torch
import numpy as np
from utils import ConvBlock, BasicBlock, Bottleneck
from fpn import FPN
from anchors import RPN
from losses import FocalLoss

import math

import sys, os

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}

def layers_dim(h_x, w_x, factor, n_layers):
    h = []
    w = []
    for i in range(n_layers):
        h.append(h_x)
        w.append(w_x)
        h_x /= float(factor)
        w_x /= float(factor)
    return h, w


    
class RegressionModel(nn.Module):
    """
    Input:
        - in_plane : features generated by FPN model at a given level.
    """
    
    def __init__(self, in_plane, n_layers=4, num_anchors=9, feature_size=256):
        super(RegressionModel, self).__init__()
        
        in_features = [in_plane] + [feature_size] * (n_layers-1)
        
        self.convgroup = self._make_convgroup(n_layers, in_features, feature_size)
        self.output = nn.Conv2d(feature_size, num_anchors * 4, kernel_size=3, stride=1, padding=1)
        

    def _make_convgroup(self, n_layers, in_features, out_features):
        layers = []
        for i in range(n_layers):
            layers.append( nn.Conv2d(in_features[i], out_features, kernel_size=3, padding=1) )
            layers.append( nn.ReLU() )
            
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.convgroup(x)
        out = self.output(out)
        
        # out is B x C x W x H, with C = 4*num_anchors
        out = out.permute(0, 2, 3, 1)
        
        out = out.contiguous().view(out.shape[0], -1, 4)
        
        return out
        
        

class ClassificationModel(nn.Module):
    """ Uses sigmoid, but could be extended to CE. """
    def __init__(self, in_plane, n_layers=4, num_anchors=9, num_classes=80, prior=0.01, feature_size=256):
        super(ClassificationModel, self).__init__()
        
        in_features = [in_plane] + [feature_size] * (n_layers-1)
        self.num_anchors = num_anchors
        self.num_classes = num_classes
        
        self.convgroup = self._make_convgroup(n_layers, in_features, feature_size)
        self.output = nn.Conv2d(feature_size, num_anchors * num_classes, kernel_size=3, stride=1, padding=1)
        self.output_act = nn.Sigmoid()
        
        
    def _make_convgroup(self, n_layers, in_features, out_features):
        layers = []
        for i in range(n_layers):
            layers.append( nn.Conv2d(in_features[i], out_features, kernel_size=3, padding=1) )
            layers.append( nn.ReLU() )
            
        return nn.Sequential(*layers)
        
    
    def forward(self, x):
        out = self.convgroup(x)
        out = self.output(out)
        out = self.output_act(out)
        # Output: [N,C,H,W], with C = num_classes * num_anchors
        
        out1 = out.permute(0, 2, 3, 1)
        N, H, W, C = out1.size()
        out2 = out1.view(N, H, W, self.num_anchors, self.num_classes)
        
        # [N, HxWxnum_anchors, num_classes]
        out3 = out2.contiguous().view(x.size(0), -1, self.num_classes)
        
        return out3
        
    
    
class RetinaNet(nn.Module):
    """ 
    Ideally, it should extend to allow more multiple type of backbones.
    """
    
    def __init__(self,):
        super(RetinaNet, self).__init__()
        
    def forward(self):
        pass
        

        
class ResNet(nn.Module):
    
    def __init__(self, layers, block, num_classes):
        super(ResNet, self).__init__()
        
        self.training = True
        
        self.in_plane = 64
        
        self.conv1 = ConvBlock(3, 64, kernel=7, stride=2, pad=3, bias=False)
        
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        # Needs to be converted toa for loop
        self.conv_2 = self._make_layers(block, layers[0], 64)        
        self.conv_3 = self._make_layers(block, layers[1], 128, stride=2)
        self.conv_4 = self._make_layers(block, layers[2], 256, stride=2)
        self.conv_5 = self._make_layers(block, layers[3], 512, stride=2)
        
        # Feature Pyramid Network
        self.fpn = FPN([128, 256, 512])
        
        # Regression Model
        self.regressionModel = RegressionModel(256)
        self.classificationModel = ClassificationModel(256, num_classes=num_classes)
        
        # Anchors
        self.anchors = RPN()
        
        # Focal Loss
        self.focalLoss = FocalLoss()
        
        
        # Utils Function
        self.regressBoxes = BBoxTransform()
        self.clipBoxes = ClipBoxes()
        
        self._init_weights()
        
    
    def _init_weights(self):
        conv_count = 0
        bn_count = 0
        for module in self.modules():
            if isinstance(module, nn.Conv2d):
                # print(module)
                n = module.kernel_size[0] * module.kernel_size[1] * module.out_channels
                module.weight.data.normal_(0, math.sqrt(2. / n))
                conv_count += 1
            elif isinstance(module, nn.BatchNorm2d):
                # print(module)
                module.weight.data.fill_(1)
                module.bias.data.zero_()
                bn_count += 1
                
        prior = 0.01
        
        self.classificationModel.output.weight.data.fill_(0)
        self.classificationModel.output.bias.data.fill_(-np.log( (1-prior) / prior) )
        
        self.regressionModel.output.weight.data.fill_(0)
        self.regressionModel.output.bias.data.fill_(0)
        
        self.freeze_bn()
                
        print("Initialized Conv Layers:",conv_count)
        print("Initialized BatchNorm Layers:",bn_count)
        
    
    def freeze_bn(self):
        ''' Freeze BatchNorm Layers '''
        for module in self.modules():
            if isinstance(module, nn.BatchNorm2d):
                for param in module.parameters():
                    param.requires_grad = False
        
    def _make_layers(self, block, n_layers, out_plane, stride=1):
        
        downsample = None
        
        # Assume downsampling is performed by adding an extra layer to perform downsampling
        if self.in_plane != out_plane or block == Bottleneck:
            downsample = ConvBlock(self.in_plane, out_plane * block.expansion, kernel=1, stride=stride, pad=0, bias=False, is_relu=False)
        
        layers = [ block(self.in_plane, out_plane, stride, downsample) ]
        self.in_plane = out_plane * block.expansion
        for i in range(1, n_layers):
            layers.append( block(self.in_plane, out_plane) )
            
        return nn.Sequential(*layers)
        
        
    def forward(self, inputs):
        
        if self.training:
            img_batch, annotations = inputs
        
        x = self.conv1(img_batch)
        x = self.maxpool(x)
        
        # This could be for-loop
        x1 = self.conv_2(x)
        x2 = self.conv_3(x1)
        x3 = self.conv_4(x2)
        x4 = self.conv_5(x3)
        
        # Reverse order for top-to-bottom pass
        features = self.fpn([x4,x3,x2])
        
        # Regression-Box (Different): list of varying sizes
        regression = torch.cat([self.regressionModel(feature) for feature in features], dim=1)

        # Classification
        classification = torch.cat([self.classificationModel(feature) for feature in features], dim=1)
        
        # Anchors
        anchors = self.anchors(img_batch)
        
        if self.training:
            return self.focalLoss(classification, regression, anchors, annotations)
        
        
        
        return None
    
# ==================== LOAD MODEL ======================
    
def resnet18(num_classes, pretrained=False, **kwargs):
    """ Constructs ResNet Model """
    model = ResNet([2, 2, 2, 2], BasicBlock, num_classes)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet18'], model_dir='models'), strict=False)
    return model

def resnet34(num_classes, pretrained=False, **kwargs):
    """ Constructs ResNet Model """
    model = ResNet([3, 4, 6, 3], BasicBlock, num_classes)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet34'], model_dir='models'), strict=False)
    return model

def resnet50(num_classes, pretrained=False, **kwargs):
    """ Constructs ResNet Model """
    model = ResNet([3, 4, 6, 3], Bottleneck, num_classes)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet50'], model_dir='models'), strict=False)
    return model

def resnet101(num_classes, pretrained=False, **kwargs):
    """ Constructs ResNet Model """
    model = ResNet([3, 4, 23, 3], Bottleneck, num_classes)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet101'], model_dir='models'), strict=False)
    return model

def resnet152(num_classes, pretrained=False, **kwargs):
    """ Constructs ResNet Model """
    model = ResNet([3, 8, 36, 3], Bottleneck, num_classes)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet152'], model_dir='models'), strict=False)
    return model
    
    
    
if __name__ == '__main__':
    print(">>> Start...")
    
    os.environ["CUDA_VISIBLE_DEVICES"]="2"
    
    x = torch.randn((10, 3, 224, 224)).cuda()
    x_dim = np.array([x.size(2), x.size(3)])
    layers = [3, 4, 6, 3]
    
    model = ResNet(layers, Bottleneck, 80).cuda()
    model.forward((x, None))