Upload 5 files

nets/CSPdarknet_tiny.py

@@ -0,0 +1,101 @@
+import math
+import cv2
+import torch
+import torch.nn as nn
+import torchvision
+from nets.FAENet import FAENet
+from torch import nn
+import torch
+
+class BasicConv(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
+        super(BasicConv, self).__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size // 2, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.activation = nn.LeakyReLU(0.1)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.activation(x)
+        return x
+    
+class Resblock_body(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Resblock_body, self).__init__()
+        self.out_channels = out_channels
+
+        self.conv1 = BasicConv(in_channels, out_channels, 3)
+        self.conv2 = BasicConv(out_channels // 2, out_channels // 2, 3)
+        self.conv3 = BasicConv(out_channels // 2, out_channels // 2, 3)
+        self.conv4 = BasicConv(out_channels, out_channels, 1)
+        self.maxpool = nn.MaxPool2d([2, 2], [2, 2])
+
+    def forward(self, x):
+        x = self.conv1(x)
+        route = x
+        c = self.out_channels
+        x = torch.split(x, c // 2, dim=1)[1]
+        x = self.conv2(x)
+        route1 = x
+        x = self.conv3(x)
+        x = torch.cat([x, route1], dim=1)
+        x = self.conv4(x)
+        feat = x
+        x = torch.cat([route, x], dim=1)
+        x = self.maxpool(x)
+        return x, feat
+
+class CSPDarkNet(nn.Module):
+    def __init__(self):
+        super(CSPDarkNet, self).__init__()
+
+        self.faenet = FAENet()
+        # 416,416,3 -> 208,208,32 -> 104,104,64
+        self.conv1 = BasicConv(3, 32, kernel_size=3, stride=2)
+        self.conv2 = BasicConv(32, 64, kernel_size=3, stride=2)
+        # 104,104,64 -> 52,52,128
+        self.resblock_body1 = Resblock_body(64, 64)
+        # 52,52,128 -> 26,26,256
+        self.resblock_body2 = Resblock_body(128, 128)
+        # 26,26,256 -> 13,13,512
+        self.resblock_body3 = Resblock_body(256, 256)
+        # 13,13,512 -> 13,13,512
+        self.conv3 = BasicConv(512, 512, kernel_size=3)
+        self.num_features = 1
+        # 进行权值初始化
+        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 forward(self, x):
+
+        x = self.faenet(x)
+        # 416,416,3 -> 208,208,32 -> 104,104,64
+        x = self.conv1(x)
+        x = self.conv2(x)
+        # 104,104,64 -> 52,52,128
+        x, _ = self.resblock_body1(x)
+        # 52,52,128 -> 26,26,256
+        x, _ = self.resblock_body2(x)
+        # 26,26,256 -> 13,13,512
+        #           -> feat1 26,26,256
+        x, feat1 = self.resblock_body3(x)
+        # 13,13,512 -> 13,13,512
+        x = self.conv3(x)
+        feat2 = x
+
+        return feat1, feat2
+
+
+def darknet_tiny(pretrained, **kwargs):
+    model = CSPDarkNet()
+    if pretrained:
+        model.load_state_dict(torch.load("model_data/CSPdarknet53_tiny_backbone_weights.pth"))
+    return model
+

nets/FAENet.py

@@ -0,0 +1,248 @@
+import math
+import cv2
+import torch
+import torch.nn as nn
+import torchvision
+from torch import nn
+import torch
+
+class eca_block(nn.Module):
+    def __init__(self, channel, b=1, gamma=2):
+        super(eca_block, self).__init__()
+        kernel_size = int(abs((math.log(channel, 2) + b) / gamma))
+        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
+        
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False) 
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+
+        y = self.avg_pool(x)
+        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
+        y = self.sigmoid(y)
+        return x * y.expand_as(x)
+    
+class DilatedConvNet(nn.Module):
+    def __init__(self, in_channels, out_channels, dilation, padding, kernel_size):
+        super(DilatedConvNet, self).__init__()
+        self.dilated_conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, dilation=dilation)
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+
+        x = self.dilated_conv(x)
+        x = self.relu(x)
+
+        return x
+
+class LAM(nn.Module):
+    def __init__(self, ch=16):
+        super().__init__()
+        self.eca = eca_block(ch)
+        self.conv1 = nn.Conv2d(6, 3, 3, padding=1)
+
+    def forward(self, x):
+        x = self.eca(x)
+        x = self.conv1(x)
+        return x
+
+class RFEM(nn.Module):
+    def __init__(
+            self,
+            ch_blocks=64,
+            ch_mask=16,
+    ):
+        super().__init__()
+
+        self.encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1),
+                                     nn.LeakyReLU(True),
+                                     nn.Conv2d(16, ch_blocks, 3, padding=1),
+                                     nn.LeakyReLU(True))
+
+        self.dconv1 = DilatedConvNet(ch_blocks,
+                                  ch_blocks // 4,
+                                  kernel_size=3,
+                                  padding=1, dilation=1)
+        self.dconv2 = DilatedConvNet(ch_blocks,
+                                  ch_blocks // 4,
+                                  kernel_size=3,
+                                  padding=2, dilation=2)
+        self.dconv3 = DilatedConvNet(ch_blocks,
+                                  ch_blocks // 4,
+                                  kernel_size=3,
+                                  padding=3, dilation=3)
+        self.dconv4 = nn.Conv2d(ch_blocks,
+                                  ch_blocks // 4,
+                                  kernel_size=7,
+                                  padding=3)
+
+        self.decoder = nn.Sequential(nn.Conv2d(ch_blocks, 16, 3, padding=1),
+                                     nn.LeakyReLU(True),
+                                     nn.Conv2d(16, 3, 3, padding=1),
+                                     nn.LeakyReLU(True),
+                                     )
+
+        self.lam = LAM(ch_mask)
+
+    def forward(self, x):
+        x1 = self.encoder(x)
+        x1_1 = self.dconv1(x1)
+        x1_2 = self.dconv2(x1)
+        x1_3 = self.dconv3(x1)
+        x1_4 = self.dconv4(x1)
+        x1 = torch.cat([x1_1, x1_2, x1_3, x1_4], dim=1)
+        x1 = self.decoder(x1)
+        out = x + x1
+        out = torch.relu(out)
+        mask = self.lam(torch.cat([x, out], dim=1))
+        return out, mask
+
+class ATEM(nn.Module):
+    def __init__(self, in_ch=3, inter_ch=32, out_ch=3, kernel_size=3):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(in_ch, inter_ch, kernel_size, padding=kernel_size // 2),
+            nn.LeakyReLU(True),
+        )
+        self.shift_conv = nn.Sequential(
+            nn.Conv2d(in_ch, inter_ch, kernel_size, padding=kernel_size // 2))
+        self.scale_conv = nn.Sequential(
+            nn.Conv2d(in_ch, inter_ch, kernel_size, padding=kernel_size // 2))
+
+
+        self.decoder = nn.Sequential(
+            nn.Conv2d(inter_ch, out_ch, kernel_size, padding=kernel_size // 2))
+
+    def forward(self, x, tag):
+        x = self.encoder(x)
+        scale = self.scale_conv(tag)
+        shift = self.shift_conv(tag)
+        x = x +(x * scale + shift)
+        x = self.decoder(x)
+        return x
+
+class Trans_high(nn.Module):
+    def __init__(self, in_ch=3, inter_ch=16, out_ch=3, kernel_size=3):
+        super().__init__()
+        self.atem = ATEM(in_ch, inter_ch, out_ch, kernel_size)
+    def forward(self, x, tag):
+        x = x + self.atem(x, tag)
+        return x
+
+
+class Up_tag(nn.Module):
+    def __init__(self, kernel_size=1, ch=3):
+        super().__init__()
+        self.up = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+            nn.Conv2d(ch,
+                      ch,
+                      kernel_size,
+                      stride=1,
+                      padding=kernel_size // 2,
+                      bias=False))
+
+    def forward(self, x):
+        x = self.up(x)
+        return x
+
+
+class Lap_Pyramid_Conv(nn.Module):
+    def __init__(self, num_high=3, kernel_size=5, channels=3):
+        super().__init__()
+
+        self.num_high = num_high
+        self.kernel = self.gauss_kernel(kernel_size, channels)
+
+    def gauss_kernel(self, kernel_size, channels):
+        kernel = cv2.getGaussianKernel(kernel_size, 0).dot(
+            cv2.getGaussianKernel(kernel_size, 0).T)
+        kernel = torch.FloatTensor(kernel).unsqueeze(0).repeat(
+            channels, 1, 1, 1)
+        kernel = torch.nn.Parameter(data=kernel, requires_grad=False)
+        return kernel
+
+    def conv_gauss(self, x, kernel):
+        n_channels, _, kw, kh = kernel.shape
+        x = torch.nn.functional.pad(x, (kw // 2, kh // 2, kw // 2, kh // 2),
+                                mode='reflect') 
+        x = torch.nn.functional.conv2d(x, kernel, groups=n_channels)
+        return x
+    def downsample(self, x):
+        return x[:, :, ::2, ::2]
+    def pyramid_down(self, x):
+        return self.downsample(self.conv_gauss(x, self.kernel))
+    def upsample(self, x):
+        up = torch.zeros((x.size(0), x.size(1), x.size(2) * 2, x.size(3) * 2),
+                         device=x.device)
+        up[:, :, ::2, ::2] = x * 4
+
+        return self.conv_gauss(up, self.kernel)
+
+    def pyramid_decom(self, img):
+        self.kernel = self.kernel.to(img.device)
+        current = img
+        pyr = []
+        for _ in range(self.num_high):
+            down = self.pyramid_down(current)
+            up = self.upsample(down)
+            diff = current - up
+            pyr.append(diff)
+            current = down
+        pyr.append(current)
+        return pyr
+
+    def pyramid_recons(self, pyr):
+        image = pyr[0]
+        for level in pyr[1:]:
+            up = self.upsample(image)
+            image = up + level
+        return image
+
+class FAENet(nn.Module):
+    def __init__(self,
+                 num_high=1,
+                 ch_blocks=32,
+                 up_ksize=1,
+                 high_ch=32,
+                 high_ksize=3,
+                 ch_mask=32,
+                 gauss_kernel=7):
+        super().__init__()
+        self.num_high = num_high
+        self.lap_pyramid = Lap_Pyramid_Conv(num_high, gauss_kernel)
+        self.rfem = RFEM(ch_blocks, ch_mask)
+
+        for i in range(0, self.num_high):
+            self.__setattr__('up_tag_layer_{}'.format(i),
+                             Up_tag(up_ksize, ch=3))
+            self.__setattr__('trans_high_layer_{}'.format(i),
+                             Trans_high(3, high_ch, 3, high_ksize))
+
+    def forward(self, x):
+        pyrs = self.lap_pyramid.pyramid_decom(img=x)
+
+        trans_pyrs = []
+        trans_pyr, tag = self.rfem(pyrs[-1])
+        trans_pyrs.append(trans_pyr)
+
+        commom_tag = []
+        for i in range(self.num_high):
+            tag = self.__getattr__('up_tag_layer_{}'.format(i))(tag)
+            commom_tag.append(tag)
+
+        for i in range(self.num_high):
+            trans_pyr = self.__getattr__('trans_high_layer_{}'.format(i))(
+                pyrs[-2 - i], commom_tag[i])
+            trans_pyrs.append(trans_pyr)
+
+        out = self.lap_pyramid.pyramid_recons(trans_pyrs)
+
+        return out
+    
+
+faenet = FAENet()
+params = faenet.parameters()
+num_params = sum(p.numel() for p in params)
+print("FAENet parameters: {:.2f}K   ".format(num_params/ 1024) + "{:.2f} MB".format(num_params/ (1024 * 1024)))

nets/__init__.py

@@ -0,0 +1 @@
+#

nets/yolo.py

@@ -0,0 +1,241 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from nets.CSPdarknet import darknet
+from nets.CSPdarknet_tiny import darknet_tiny
+from nets.mobilenetv2 import mobilenet_v2
+from nets.shufflenet_v2 import shufflenet_v2
+from nets.ghostnet import ghostnet
+from nets.attention import cbam_block, eca_block, se_block, CA_Block
+attention_block = [se_block, cbam_block, eca_block, CA_Block]
+
+
+
+#-------------------------------------------------#
+#   卷积块 -> 卷积 + 标准化 + 激活函数
+#   Conv2d + BatchNormalization + LeakyReLU
+#-------------------------------------------------#
+class BasicConv(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
+        super(BasicConv, self).__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.activation = nn.LeakyReLU(0.1)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.activation(x)
+        return x
+
+#---------------------------------------------------#
+#   卷积 + 上采样
+#---------------------------------------------------#
+class Upsample(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Upsample, self).__init__()
+
+        self.upsample = nn.Sequential(
+            BasicConv(in_channels, out_channels, 1),
+            nn.Upsample(scale_factor=2, mode='nearest')
+        )
+
+    def forward(self, x,):
+        x = self.upsample(x)
+        return x
+
+#---------------------------------------------------#
+#   最后获得yolov4的输出
+#---------------------------------------------------#
+def yolo_head(filters_list, in_filters):
+    m = nn.Sequential(
+        BasicConv(in_filters, filters_list[0], 3),
+        nn.Conv2d(filters_list[0], filters_list[1], 1),
+    )
+    return m
+
+class ConvBNReLU(nn.Module):
+    '''Module for the Conv-BN-ReLU tuple.'''
+
+    def __init__(self, c_in, c_out, kernel_size, stride, padding, dilation,
+                 use_relu=True):
+        super(ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(
+            c_in, c_out, kernel_size=kernel_size, stride=stride,
+            padding=padding, dilation=dilation, bias=False)
+        self.bn = nn.BatchNorm2d(c_out)
+        if use_relu:
+            self.relu = nn.ReLU(inplace=True)
+        else:
+            self.relu = None
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.relu is not None:
+            x = self.relu(x)
+        return x
+
+
+class CARAFE(nn.Module):
+    def __init__(self, c, c_mid=64, scale=2, k_up=5, k_enc=3):
+        """ The unofficial implementation of the CARAFE module.
+
+        The details are in "https://arxiv.org/abs/1905.02188".
+
+        Args:
+            c: The channel number of the input and the output.
+            c_mid: The channel number after compression.
+            scale: The expected upsample scale.
+            k_up: The size of the reassembly kernel.
+            k_enc: The kernel size of the encoder.
+
+        Returns:
+            X: The upsampled feature map.
+        """
+        super(CARAFE, self).__init__()
+        self.scale = scale
+
+        self.comp = ConvBNReLU(c, c_mid, kernel_size=1, stride=1,
+                               padding=0, dilation=1)
+        self.enc = ConvBNReLU(c_mid, (scale * k_up) ** 2, kernel_size=k_enc,
+                              stride=1, padding=k_enc // 2, dilation=1,
+                              use_relu=False)
+        self.pix_shf = nn.PixelShuffle(scale)
+
+        self.upsmp = nn.Upsample(scale_factor=scale, mode='nearest')
+        self.unfold = nn.Unfold(kernel_size=k_up, dilation=scale,
+                                padding=k_up // 2 * scale)
+
+    def forward(self, X):
+        b, c, h, w = X.size()
+        h_, w_ = h * self.scale, w * self.scale
+
+        W = self.comp(X)  # b * m * h * w
+        W = self.enc(W)  # b * 100 * h * w
+        W = self.pix_shf(W)  # b * 25 * h_ * w_
+        W = F.softmax(W, dim=1)  # b * 25 * h_ * w_
+
+        X = self.upsmp(X)  # b * c * h_ * w_
+        X = self.unfold(X)  # b * 25c * h_ * w_
+        X = X.view(b, c, -1, h_, w_)  # b * 25 * c * h_ * w_
+
+        X = torch.einsum('bkhw,bckhw->bchw', [W, X])  # b * c * h_ * w_
+        return X
+
+#---------------------------------------------------#
+#   yolo_body--MSFNet
+#---------------------------------------------------#
+class YoloBody(nn.Module):
+    def __init__(self, anchors_mask, num_classes, phi=0, backbone ='', pretrained=False):
+        super(YoloBody, self).__init__()
+
+        self.phi            = phi
+
+        if backbone == 'cspdarknet':
+            self.backbone  = darknet(pretrained)
+            self.conv_for_P5    = BasicConv(512,256,1)
+            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
+            self.upsample_1       = Upsample(256,128)
+            self.conv1 = BasicConv(256,128,1)
+            self.upsample_2 = CARAFE(128)
+            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],384)
+
+            if 1 <= self.phi and self.phi <= 4:
+                self.feat1_att      = attention_block[self.phi - 1](256)
+                self.feat2_att      = attention_block[self.phi - 1](512)
+                self.upsample_att   = attention_block[self.phi - 1](128)  
+        
+        elif backbone == 'tiny':
+            self.backbone  = darknet_tiny(pretrained)
+            self.conv_for_P5    = BasicConv(512,256,1)
+            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
+            self.upsample_1       = Upsample(256,128)
+            self.conv1 = BasicConv(256,128,1)
+            self.upsample_2 = CARAFE(128)
+            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],384)
+
+            if 1 <= self.phi and self.phi <= 4:
+                self.feat1_att      = attention_block[self.phi - 1](256)
+                self.feat2_att      = attention_block[self.phi - 1](512)
+                self.upsample_att   = attention_block[self.phi - 1](128)  
+
+        elif backbone == 'mobilenetv2':
+            self.backbone  = mobilenet_v2(pretrained)
+            self.conv_for_P5    = BasicConv(320,256,1)
+            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
+            self.upsample_1       = Upsample(256,128)
+            self.conv1 = BasicConv(256,128,1)
+            self.upsample_2 = CARAFE(128)
+            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],224)
+
+            if 1 <= self.phi and self.phi <= 4:
+                self.feat1_att      = attention_block[self.phi - 1](256)
+                self.feat2_att      = attention_block[self.phi - 1](512)
+                self.upsample_att   = attention_block[self.phi - 1](128)  
+
+        elif backbone == 'shufflenetv2':
+            self.backbone  = shufflenet_v2()
+            self.conv_for_P5    = BasicConv(1024,256,1)
+            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
+            self.upsample_1       = Upsample(256,128)
+            self.conv1 = BasicConv(256,128,1)
+            self.upsample_2 = CARAFE(128)
+            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],592)
+
+            if 1 <= self.phi and self.phi <= 4:
+                self.feat1_att      = attention_block[self.phi - 1](256)
+                self.feat2_att      = attention_block[self.phi - 1](512)
+                self.upsample_att   = attention_block[self.phi - 1](128)  
+
+        elif backbone == 'ghostnet':
+            self.backbone  = ghostnet()
+            self.conv_for_P5    = BasicConv(160,256,1)
+            self.yolo_headP5    = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
+            self.upsample_1       = Upsample(256,128)
+            self.conv1 = BasicConv(256,128,1)
+            self.upsample_2 = CARAFE(128)
+            self.yolo_headP4    = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],240)
+
+            if 1 <= self.phi and self.phi <= 4:
+                self.feat1_att      = attention_block[self.phi - 1](256)
+                self.feat2_att      = attention_block[self.phi - 1](512)
+                self.upsample_att   = attention_block[self.phi - 1](128)  
+
+        
+
+    def forward(self, x):
+
+        #---------------------------------------------------#
+        #   生成CSPdarknet53_tiny的主干模型
+        #   feat1的shape为26,26,256
+        #   feat2的shape为13,13,512
+        #---------------------------------------------------#
+        feat1, feat2 = self.backbone(x)
+        if 1 <= self.phi and self.phi <= 4:
+            feat1 = self.feat1_att(feat1)
+            feat2 = self.feat2_att(feat2)
+
+        # 13,13,512 -> 13,13,256
+        P5 = self.conv_for_P5(feat2)
+        # 13,13,256 -> 13,13,512 -> 13,13,255
+        out0 = self.yolo_headP5(P5)
+
+        P6 = self.conv_for_P5(feat2)
+        P6_Upsample = self.upsample_1(P6)
+
+        # 13,13,256 -> 13,13,128 -> 26,26,128  卷积+轻量级上采样
+        P5 = self.conv1(P5)
+        P5_Upsample = self.upsample_2(P5)
+
+        sum = P5_Upsample + P6_Upsample
+        # 26,26,256 + 26,26,128 -> 26,26,384
+        # if 1 <= self.phi and self.phi <= 4:
+        #     P5_Upsample = self.upsample_att(P5_Upsample)
+        # P4 = torch.cat([P5_Upsample, feat1],axis=1)
+        P4 = torch.cat([sum, feat1],axis=1)
+        # 26,26,384 -> 26,26,256 -> 26,26,255
+        out1 = self.yolo_headP4(P4)
+        
+        return out0, out1

nets/yolo_training.py

@@ -0,0 +1,277 @@
+import math
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+class YOLOLoss(nn.Module):
+    def __init__(self, anchors, num_classes, input_shape, cuda, anchors_mask = [[6,7,8], [3,4,5], [0,1,2]], label_smoothing = 0):
+        super(YOLOLoss, self).__init__()
+        self.anchors        = anchors
+        self.num_classes    = num_classes
+        self.bbox_attrs     = 5 + num_classes
+        self.input_shape    = input_shape
+        self.anchors_mask   = anchors_mask
+        self.label_smoothing = label_smoothing
+        self.balance        = [0.4, 1.0, 4]
+        self.box_ratio      = 0.05
+        self.obj_ratio      = 5 * (input_shape[0] * input_shape[1]) / (416 ** 2)
+        self.cls_ratio      = 1 * (num_classes / 80)
+
+        self.ignore_threshold = 0.5
+        self.cuda           = cuda
+
+    def clip_by_tensor(self, t, t_min, t_max):
+        t = t.float()
+        result = (t >= t_min).float() * t + (t < t_min).float() * t_min
+        result = (result <= t_max).float() * result + (result > t_max).float() * t_max
+        return result
+
+    def MSELoss(self, pred, target):
+        return torch.pow(pred - target, 2)
+
+    def BCELoss(self, pred, target):
+        epsilon = 1e-7
+        pred    = self.clip_by_tensor(pred, epsilon, 1.0 - epsilon)
+        output  = - target * torch.log(pred) - (1.0 - target) * torch.log(1.0 - pred)
+        return output
+        
+    def box_ciou(self, b1, b2):
+        b1_xy       = b1[..., :2]
+        b1_wh       = b1[..., 2:4]
+        b1_wh_half  = b1_wh/2.
+        b1_mins     = b1_xy - b1_wh_half
+        b1_maxes    = b1_xy + b1_wh_half
+        b2_xy       = b2[..., :2]
+        b2_wh       = b2[..., 2:4]
+        b2_wh_half  = b2_wh/2.
+        b2_mins     = b2_xy - b2_wh_half
+        b2_maxes    = b2_xy + b2_wh_half
+
+        intersect_mins  = torch.max(b1_mins, b2_mins)
+        intersect_maxes = torch.min(b1_maxes, b2_maxes)
+        intersect_wh    = torch.max(intersect_maxes - intersect_mins, torch.zeros_like(intersect_maxes))
+        intersect_area  = intersect_wh[..., 0] * intersect_wh[..., 1]
+        b1_area         = b1_wh[..., 0] * b1_wh[..., 1]
+        b2_area         = b2_wh[..., 0] * b2_wh[..., 1]
+        union_area      = b1_area + b2_area - intersect_area
+        iou             = intersect_area / torch.clamp(union_area,min = 1e-6)
+
+        center_distance = torch.sum(torch.pow((b1_xy - b2_xy), 2), axis=-1)
+        
+        enclose_mins    = torch.min(b1_mins, b2_mins)
+        enclose_maxes   = torch.max(b1_maxes, b2_maxes)
+        enclose_wh      = torch.max(enclose_maxes - enclose_mins, torch.zeros_like(intersect_maxes))
+        enclose_diagonal = torch.sum(torch.pow(enclose_wh,2), axis=-1)
+        ciou            = iou - 1.0 * (center_distance) / torch.clamp(enclose_diagonal,min = 1e-6)
+        
+        v       = (4 / (math.pi ** 2)) * torch.pow((torch.atan(b1_wh[..., 0] / torch.clamp(b1_wh[..., 1],min = 1e-6)) - torch.atan(b2_wh[..., 0] / torch.clamp(b2_wh[..., 1], min = 1e-6))), 2)
+        alpha   = v / torch.clamp((1.0 - iou + v), min=1e-6)
+        ciou    = ciou - alpha * v
+        return ciou
+
+    def smooth_labels(self, y_true, label_smoothing, num_classes):
+        return y_true * (1.0 - label_smoothing) + label_smoothing / num_classes
+
+    def forward(self, l, input, targets=None):
+        bs      = input.size(0)
+        in_h    = input.size(2)
+        in_w    = input.size(3)
+        stride_h = self.input_shape[0] / in_h
+        stride_w = self.input_shape[1] / in_w
+        scaled_anchors  = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in self.anchors]
+
+        prediction = input.view(bs, len(self.anchors_mask[l]), self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
+        
+        x = torch.sigmoid(prediction[..., 0])
+        y = torch.sigmoid(prediction[..., 1])
+        w = prediction[..., 2]
+        h = prediction[..., 3]
+        conf = torch.sigmoid(prediction[..., 4])
+        pred_cls = torch.sigmoid(prediction[..., 5:])
+
+        y_true, noobj_mask, box_loss_scale = self.get_target(l, targets, scaled_anchors, in_h, in_w)
+
+        noobj_mask, pred_boxes = self.get_ignore(l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask)
+
+        if self.cuda:
+            y_true          = y_true.type_as(x)
+            noobj_mask      = noobj_mask.type_as(x)
+            box_loss_scale  = box_loss_scale.type_as(x)
+        box_loss_scale = 2 - box_loss_scale
+
+        loss        = 0
+        obj_mask    = y_true[..., 4] == 1
+        n           = torch.sum(obj_mask)
+        if n != 0:
+            ciou        = self.box_ciou(pred_boxes, y_true[..., :4]).type_as(x)
+            loss_loc    = torch.mean((1 - ciou)[obj_mask])
+            
+            loss_cls    = torch.mean(self.BCELoss(pred_cls[obj_mask], y_true[..., 5:][obj_mask]))
+            loss        += loss_loc * self.box_ratio + loss_cls * self.cls_ratio
+
+        loss_conf   = torch.mean(self.BCELoss(conf, obj_mask.type_as(conf))[noobj_mask.bool() | obj_mask])
+        loss        += loss_conf * self.balance[l] * self.obj_ratio
+        return loss
+
+    def calculate_iou(self, _box_a, _box_b):
+        b1_x1, b1_x2 = _box_a[:, 0] - _box_a[:, 2] / 2, _box_a[:, 0] + _box_a[:, 2] / 2
+        b1_y1, b1_y2 = _box_a[:, 1] - _box_a[:, 3] / 2, _box_a[:, 1] + _box_a[:, 3] / 2
+        b2_x1, b2_x2 = _box_b[:, 0] - _box_b[:, 2] / 2, _box_b[:, 0] + _box_b[:, 2] / 2
+        b2_y1, b2_y2 = _box_b[:, 1] - _box_b[:, 3] / 2, _box_b[:, 1] + _box_b[:, 3] / 2
+
+        box_a = torch.zeros_like(_box_a)
+        box_b = torch.zeros_like(_box_b)
+        box_a[:, 0], box_a[:, 1], box_a[:, 2], box_a[:, 3] = b1_x1, b1_y1, b1_x2, b1_y2
+        box_b[:, 0], box_b[:, 1], box_b[:, 2], box_b[:, 3] = b2_x1, b2_y1, b2_x2, b2_y2
+
+        A = box_a.size(0)
+        B = box_b.size(0)
+
+        max_xy  = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
+        min_xy  = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2))
+        inter   = torch.clamp((max_xy - min_xy), min=0)
+        inter   = inter[:, :, 0] * inter[:, :, 1]
+        area_a = ((box_a[:, 2]-box_a[:, 0]) * (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  
+        area_b = ((box_b[:, 2]-box_b[:, 0]) * (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  
+        union = area_a + area_b - inter
+        return inter / union 
+    
+    def get_target(self, l, targets, anchors, in_h, in_w):
+        bs              = len(targets)
+        noobj_mask      = torch.ones(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False)
+        box_loss_scale  = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False)
+        y_true          = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, self.bbox_attrs, requires_grad = False)
+        for b in range(bs):            
+            if len(targets[b])==0:
+                continue
+            batch_target = torch.zeros_like(targets[b])
+            batch_target[:, [0,2]] = targets[b][:, [0,2]] * in_w
+            batch_target[:, [1,3]] = targets[b][:, [1,3]] * in_h
+            batch_target[:, 4] = targets[b][:, 4]
+            batch_target = batch_target.cpu()
+            
+            gt_box          = torch.FloatTensor(torch.cat((torch.zeros((batch_target.size(0), 2)), batch_target[:, 2:4]), 1))
+            anchor_shapes   = torch.FloatTensor(torch.cat((torch.zeros((len(anchors), 2)), torch.FloatTensor(anchors)), 1))
+            iou     = self.calculate_iou(gt_box, anchor_shapes)
+            best_ns = torch.argmax(iou, dim=-1)
+            sort_ns = torch.argsort(iou, dim=-1, descending=True)
+
+            def check_in_anchors_mask(index, anchors_mask):
+                for sub_anchors_mask in anchors_mask:
+                    if index in sub_anchors_mask:
+                        return True
+                return False
+
+            for t, best_n in enumerate(best_ns):
+                if not check_in_anchors_mask(best_n, self.anchors_mask):
+                    for index in sort_ns[t]:
+                        if check_in_anchors_mask(index, self.anchors_mask):
+                            best_n = index
+                            break
+
+                if best_n not in self.anchors_mask[l]:
+                    continue
+                k = self.anchors_mask[l].index(best_n)
+                i = torch.floor(batch_target[t, 0]).long()
+                j = torch.floor(batch_target[t, 1]).long()
+                c = batch_target[t, 4].long()
+                
+                noobj_mask[b, k, j, i] = 0
+                y_true[b, k, j, i, 0] = batch_target[t, 0]
+                y_true[b, k, j, i, 1] = batch_target[t, 1]
+                y_true[b, k, j, i, 2] = batch_target[t, 2]
+                y_true[b, k, j, i, 3] = batch_target[t, 3]
+                y_true[b, k, j, i, 4] = 1
+                y_true[b, k, j, i, c + 5] = 1
+                box_loss_scale[b, k, j, i] = batch_target[t, 2] * batch_target[t, 3] / in_w / in_h
+        return y_true, noobj_mask, box_loss_scale
+
+    def get_ignore(self, l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask):
+        bs = len(targets)
+
+        grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_h, 1).repeat(
+            int(bs * len(self.anchors_mask[l])), 1, 1).view(x.shape).type_as(x)
+        grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_w, 1).t().repeat(
+            int(bs * len(self.anchors_mask[l])), 1, 1).view(y.shape).type_as(x)
+
+        scaled_anchors_l = np.array(scaled_anchors)[self.anchors_mask[l]]
+        anchor_w = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([0])).type_as(x)
+        anchor_h = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([1])).type_as(x)
+        
+        anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
+        anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)
+        pred_boxes_x    = torch.unsqueeze(x + grid_x, -1)
+        pred_boxes_y    = torch.unsqueeze(y + grid_y, -1)
+        pred_boxes_w    = torch.unsqueeze(torch.exp(w) * anchor_w, -1)
+        pred_boxes_h    = torch.unsqueeze(torch.exp(h) * anchor_h, -1)
+        pred_boxes      = torch.cat([pred_boxes_x, pred_boxes_y, pred_boxes_w, pred_boxes_h], dim = -1)
+        for b in range(bs):           
+            pred_boxes_for_ignore = pred_boxes[b].view(-1, 4)
+            if len(targets[b]) > 0:
+                batch_target = torch.zeros_like(targets[b])
+                batch_target[:, [0,2]] = targets[b][:, [0,2]] * in_w
+                batch_target[:, [1,3]] = targets[b][:, [1,3]] * in_h
+                batch_target = batch_target[:, :4].type_as(x)
+                anch_ious = self.calculate_iou(batch_target, pred_boxes_for_ignore)
+                anch_ious_max, _    = torch.max(anch_ious, dim = 0)
+                anch_ious_max       = anch_ious_max.view(pred_boxes[b].size()[:3])
+                noobj_mask[b][anch_ious_max > self.ignore_threshold] = 0
+        return noobj_mask, pred_boxes
+
+def weights_init(net, init_type='normal', init_gain = 0.02):
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and classname.find('Conv') != -1:
+            if init_type == 'normal':
+                torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == 'xavier':
+                torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == 'kaiming':
+                torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                torch.nn.init.orthogonal_(m.weight.data, gain=init_gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+        elif classname.find('BatchNorm2d') != -1:
+            torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
+            torch.nn.init.constant_(m.bias.data, 0.0)
+    print('initialize network with %s type' % init_type)
+    net.apply(init_func)
+
+def get_lr_scheduler(lr_decay_type, lr, min_lr, total_iters, warmup_iters_ratio = 0.05, warmup_lr_ratio = 0.1, no_aug_iter_ratio = 0.05, step_num = 10):
+    def yolox_warm_cos_lr(lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter, iters):
+        if iters <= warmup_total_iters:
+            lr = (lr - warmup_lr_start) * pow(iters / float(warmup_total_iters), 2) + warmup_lr_start
+        elif iters >= total_iters - no_aug_iter:
+            lr = min_lr
+        else:
+            lr = min_lr + 0.5 * (lr - min_lr) * (
+                1.0 + math.cos(math.pi* (iters - warmup_total_iters) / (total_iters - warmup_total_iters - no_aug_iter))
+            )
+        return lr
+
+    def step_lr(lr, decay_rate, step_size, iters):
+        if step_size < 1:
+            raise ValueError("step_size must above 1.")
+        n       = iters // step_size
+        out_lr  = lr * decay_rate ** n
+        return out_lr
+
+    if lr_decay_type == "cos":
+        warmup_total_iters  = min(max(warmup_iters_ratio * total_iters, 1), 3)
+        warmup_lr_start     = max(warmup_lr_ratio * lr, 1e-6)
+        no_aug_iter         = min(max(no_aug_iter_ratio * total_iters, 1), 15)
+        func = partial(yolox_warm_cos_lr ,lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter)
+    else:
+        decay_rate  = (min_lr / lr) ** (1 / (step_num - 1))
+        step_size   = total_iters / step_num
+        func = partial(step_lr, lr, decay_rate, step_size)
+
+    return func
+
+def set_optimizer_lr(optimizer, lr_scheduler_func, epoch):
+    lr = lr_scheduler_func(epoch)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr