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