yjh0410 1 년 전
부모
커밋
1246b35f70

+ 2 - 65
models/detectors/yolov1/yolov1_head.py

@@ -1,9 +1,7 @@
 import torch
 import torch.nn as nn
-try:
-    from .yolov1_basic import Conv
-except:
-    from yolov1_basic import Conv
+
+from .yolov1_basic import Conv
 
 
 class DecoupledHead(nn.Module):
@@ -74,64 +72,3 @@ def build_head(cfg, in_dim, out_dim, num_classes=80):
     head = DecoupledHead(cfg, in_dim, out_dim, num_classes) 
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-        'reg_max': 16,
-    }
-    fpn_dims = [256, 512, 512]
-    # Head-1
-    model = build_head(cfg, 256, fpn_dims, num_classes=80)
-    x = torch.randn(1, 256, 80, 80)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-2
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 40, 40)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-2: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-2: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-3
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 20, 20)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-3: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-3: Params : {:.2f} M'.format(params / 1e6))

+ 0 - 1
models/detectors/yolov2/yolov2.py

@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-import numpy as np
 
 from utils.misc import multiclass_nms
 

+ 2 - 65
models/detectors/yolov2/yolov2_head.py

@@ -1,9 +1,7 @@
 import torch
 import torch.nn as nn
-try:
-    from .yolov2_basic import Conv
-except:
-    from yolov2_basic import Conv
+
+from .yolov2_basic import Conv
 
 
 class DecoupledHead(nn.Module):
@@ -74,64 +72,3 @@ def build_head(cfg, in_dim, out_dim, num_classes=80):
     head = DecoupledHead(cfg, in_dim, out_dim, num_classes) 
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-        'reg_max': 16,
-    }
-    fpn_dims = [256, 512, 512]
-    # Head-1
-    model = build_head(cfg, 256, fpn_dims, num_classes=80)
-    x = torch.randn(1, 256, 80, 80)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-2
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 40, 40)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-2: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-2: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-3
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 20, 20)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-3: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-3: Params : {:.2f} M'.format(params / 1e6))

+ 2 - 65
models/detectors/yolov3/yolov3_head.py

@@ -1,9 +1,7 @@
 import torch
 import torch.nn as nn
-try:
-    from .yolov3_basic import Conv
-except:
-    from yolov3_basic import Conv
+
+from .yolov3_basic import Conv
 
 
 class DecoupledHead(nn.Module):
@@ -74,64 +72,3 @@ def build_head(cfg, in_dim, out_dim, num_classes=80):
     head = DecoupledHead(cfg, in_dim, out_dim, num_classes) 
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-        'reg_max': 16,
-    }
-    fpn_dims = [256, 512, 512]
-    # Head-1
-    model = build_head(cfg, 256, fpn_dims, num_classes=80)
-    x = torch.randn(1, 256, 80, 80)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-2
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 40, 40)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-2: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-2: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-3
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 20, 20)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-3: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-3: Params : {:.2f} M'.format(params / 1e6))

+ 2 - 65
models/detectors/yolov4/yolov4_head.py

@@ -1,9 +1,7 @@
 import torch
 import torch.nn as nn
-try:
-    from .yolov4_basic import Conv
-except:
-    from yolov4_basic import Conv
+
+from .yolov4_basic import Conv
 
 
 class DecoupledHead(nn.Module):
@@ -74,64 +72,3 @@ def build_head(cfg, in_dim, out_dim, num_classes=80):
     head = DecoupledHead(cfg, in_dim, out_dim, num_classes) 
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-        'reg_max': 16,
-    }
-    fpn_dims = [256, 512, 512]
-    # Head-1
-    model = build_head(cfg, 256, fpn_dims, num_classes=80)
-    x = torch.randn(1, 256, 80, 80)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-2
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 40, 40)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-2: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-2: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-3
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 20, 20)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-3: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-3: Params : {:.2f} M'.format(params / 1e6))

+ 2 - 65
models/detectors/yolov5/yolov5_head.py

@@ -1,9 +1,7 @@
 import torch
 import torch.nn as nn
-try:
-    from .yolov5_basic import Conv
-except:
-    from yolov5_basic import Conv
+
+from .yolov5_basic import Conv
 
 
 class DecoupledHead(nn.Module):
@@ -73,64 +71,3 @@ def build_head(cfg, in_dim, out_dim, num_classes=80):
     head = DecoupledHead(cfg, in_dim, out_dim, num_classes) 
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-        'reg_max': 16,
-    }
-    fpn_dims = [256, 512, 512]
-    # Head-1
-    model = build_head(cfg, 256, fpn_dims, num_classes=80)
-    x = torch.randn(1, 256, 80, 80)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-2
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 40, 40)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-2: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-2: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-3
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 20, 20)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-3: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-3: Params : {:.2f} M'.format(params / 1e6))

+ 2 - 4
models/detectors/yolov5/yolov5_neck.py

@@ -1,9 +1,7 @@
 import torch
 import torch.nn as nn
-try:
-    from .yolov5_basic import Conv
-except:
-    from yolov5_basic import Conv
+
+from .yolov5_basic import Conv
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher

+ 2 - 65
models/detectors/yolov7/yolov7_head.py

@@ -1,9 +1,7 @@
 import torch
 import torch.nn as nn
-try:
-    from .yolov7_basic import Conv
-except:
-    from yolov7_basic import Conv
+
+from .yolov7_basic import Conv
 
 
 class DecoupledHead(nn.Module):
@@ -74,64 +72,3 @@ def build_head(cfg, in_dim, out_dim, num_classes=80):
     head = DecoupledHead(cfg, in_dim, out_dim, num_classes) 
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-        'reg_max': 16,
-    }
-    fpn_dims = [256, 512, 512]
-    # Head-1
-    model = build_head(cfg, 256, fpn_dims, num_classes=80)
-    x = torch.randn(1, 256, 80, 80)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-2
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 40, 40)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-2: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-2: Params : {:.2f} M'.format(params / 1e6))
-
-    # Head-3
-    model = build_head(cfg, 512, fpn_dims, num_classes=80)
-    x = torch.randn(1, 512, 20, 20)
-    t0 = time.time()
-    outputs = model(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('Head-3: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-3: Params : {:.2f} M'.format(params / 1e6))

+ 0 - 1
models/detectors/yolov8/yolov8_backbone.py

@@ -76,7 +76,6 @@ class Yolov8Backbone(nn.Module):
         c3 = self.layer_3(c2)
         c4 = self.layer_4(c3)
         c5 = self.layer_5(c4)
-        print(c3)
         outputs = [c3, c4, c5]
 
         return outputs

+ 1 - 37
models/detectors/yolov8/yolov8_head.py

@@ -1,10 +1,7 @@
 import torch
 import torch.nn as nn
 
-try:
-    from .yolov8_basic import Conv
-except:
-    from yolov8_basic import Conv
+from .yolov8_basic import Conv
 
 
 # Single-level Head
@@ -130,36 +127,3 @@ def build_det_head(cfg, in_dims, num_levels=3, num_classes=80, reg_max=16):
         head = MultiLevelHead(cfg, in_dims, num_levels, num_classes, reg_max)
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'head': 'decoupled_head',
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-        'reg_max': 16,
-    }
-    fpn_dims = [256, 512, 512]
-    cls_out_dim = 256
-    reg_out_dim = 64
-    # Head-1
-    model = build_det_head(cfg, fpn_dims, num_levels=3, num_classes=80, reg_max=16)
-    print(model)
-    fpn_feats = [torch.randn(1, fpn_dims[0], 80, 80), torch.randn(1, fpn_dims[1], 40, 40), torch.randn(1, fpn_dims[2], 20, 20)]
-    t0 = time.time()
-    outputs = model(fpn_feats)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(fpn_feats, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))

+ 1 - 35
models/detectors/yolov8/yolov8_neck.py

@@ -1,10 +1,7 @@
 import torch
 import torch.nn as nn
 
-try:
-    from .yolov8_basic import Conv
-except:
-    from yolov8_basic import Conv
+from .yolov8_basic import Conv
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
@@ -71,34 +68,3 @@ def build_neck(cfg, in_dim, out_dim):
         neck = SPPFBlockCSP(cfg, in_dim, out_dim, cfg['neck_expand_ratio'])
 
     return neck
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        ## Neck: SPP
-        'neck': 'sppf',
-        'neck_expand_ratio': 0.5,
-        'pooling_size': 5,
-        'neck_act': 'silu',
-        'neck_norm': 'BN',
-        'neck_depthwise': False,
-    }
-    in_dim = 512
-    out_dim = 512
-    # Head-1
-    model = build_neck(cfg, in_dim, out_dim)
-    feat = torch.randn(1, in_dim, 20, 20)
-    t0 = time.time()
-    outputs = model(feat)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(feat, ), verbose=False)
-    print('==============================')
-    print('FPN: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('FPN: Params : {:.2f} M'.format(params / 1e6))

+ 2 - 32
models/detectors/yolov8/yolov8_pafpn.py

@@ -1,10 +1,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-try:
-    from .yolov8_basic import Conv, Yolov8StageBlock
-except:
-    from yolov8_basic import Conv, Yolov8StageBlock
+
+from .yolov8_basic import Conv, Yolov8StageBlock
 
 
 # PaFPN-ELAN
@@ -119,31 +117,3 @@ def build_fpn(cfg, in_dims):
                              depthwise=cfg['fpn_depthwise']
                              )
     return fpn_net
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'fpn': 'yolov8_pafpn',
-        'fpn_act': 'silu',
-        'fpn_norm': 'BN',
-        'fpn_depthwise': False,
-        'width': 0.25,
-        'depth': 0.34,
-        'ratio': 2.0,
-    }
-    model = build_fpn(cfg, in_dims=[64, 128, 256])
-    pyramid_feats = [torch.randn(1, 64, 80, 80), torch.randn(1, 128, 40, 40), torch.randn(1, 256, 20, 20)]
-    t0 = time.time()
-    outputs = model(pyramid_feats)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    for out in outputs:
-        print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(pyramid_feats, ), verbose=False)
-    print('==============================')
-    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Params : {:.2f} M'.format(params / 1e6))

+ 1 - 36
models/detectors/yolox/yolox_head.py

@@ -1,10 +1,7 @@
 import torch
 import torch.nn as nn
 
-try:
-    from .yolox_basic import Conv
-except:
-    from  yolox_basic import Conv
+from .yolox_basic import Conv
 
 
 class DecoupledHead(nn.Module):
@@ -74,35 +71,3 @@ def build_head(cfg, in_dim, out_dim, num_classes=80):
     head = DecoupledHead(cfg, in_dim, out_dim, num_classes) 
 
     return head
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'head': 'decoupled_head',
-        'num_cls_head': 2,
-        'num_reg_head': 2,
-        'head_act': 'silu',
-        'head_norm': 'BN',
-        'head_depthwise': False,
-    }
-    fpn_dims = [256, 256, 256]
-    cls_out_dim = 256
-    reg_out_dim = 256
-    # Head-1
-    model = build_head(cfg, 256, 256, num_classes=80)
-    print(model)
-    fpn_feats = [torch.randn(1, fpn_dims[0], 80, 80), torch.randn(1, fpn_dims[1], 40, 40), torch.randn(1, fpn_dims[2], 20, 20)]
-    t0 = time.time()
-    outputs = model(fpn_feats)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    # for out in outputs:
-    #     print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(fpn_feats, ), verbose=False)
-    print('==============================')
-    print('Head-1: GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Head-1: Params : {:.2f} M'.format(params / 1e6))

+ 1 - 4
models/detectors/yolox/yolox_neck.py

@@ -1,10 +1,7 @@
 import torch
 import torch.nn as nn
 
-try:
-    from .yolox_basic import Conv
-except:
-    from yolox_basic import Conv
+from .yolox_basic import Conv
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher

+ 1 - 34
models/detectors/yolox/yolox_pafpn.py

@@ -2,10 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-try:
-    from .yolox_basic import (Conv, build_reduce_layer, build_downsample_layer, build_fpn_block)
-except:
-    from  yolox_basic import (Conv, build_reduce_layer, build_downsample_layer, build_fpn_block)
+from .yolox_basic import (Conv, build_reduce_layer, build_downsample_layer, build_fpn_block)
 
 
 # YOLO-Style PaFPN
@@ -93,33 +90,3 @@ def build_fpn(cfg, in_dims, out_dim=None):
         fpn_net = YoloxPaFPN(cfg, in_dims, out_dim)
 
     return fpn_net
-
-
-if __name__ == '__main__':
-    import time
-    from thop import profile
-    cfg = {
-        'fpn': 'yolox_pafpn',
-        'fpn_reduce_layer': 'conv',
-        'fpn_downsample_layer': 'conv',
-        'fpn_core_block': 'cspblock',
-        'fpn_act': 'silu',
-        'fpn_norm': 'BN',
-        'fpn_depthwise': False,
-        'width': 1.0,
-        'depth': 1.0,
-    }
-    model = build_fpn(cfg, in_dims=[256, 512, 1024], out_dim=256)
-    pyramid_feats = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
-    t0 = time.time()
-    outputs = model(pyramid_feats)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    for out in outputs:
-        print(out.shape)
-
-    print('==============================')
-    flops, params = profile(model, inputs=(pyramid_feats, ), verbose=False)
-    print('==============================')
-    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Params : {:.2f} M'.format(params / 1e6))