mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-14 12:17:55 +00:00
79 lines
2.7 KiB
Python
Executable File
79 lines
2.7 KiB
Python
Executable File
import os
|
|
import json
|
|
import os.path as osp
|
|
import io
|
|
import torch
|
|
import numpy as np
|
|
from mmcv import Config
|
|
|
|
from mmdet.models import build_detector
|
|
from mmcv.cnn import get_model_complexity_info
|
|
|
|
def get_flops(cfg, input_shape):
|
|
model = build_detector(
|
|
cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
|
|
if torch.cuda.is_available():
|
|
model.cuda()
|
|
model.eval()
|
|
if hasattr(model, 'forward_dummy'):
|
|
model.forward = model.forward_dummy
|
|
else:
|
|
raise NotImplementedError(
|
|
'FLOPs counter is currently not currently supported with {}'.
|
|
format(model.__class__.__name__))
|
|
buf = io.StringIO()
|
|
all_flops, params = get_model_complexity_info(model, input_shape, print_per_layer_stat=True, as_strings=False, ost=buf)
|
|
buf = buf.getvalue()
|
|
lines = buf.split("\n")
|
|
names = ['(stem)', '(layer1)', '(layer2)', '(layer3)', '(layer4)', '(neck)', '(bbox_head)']
|
|
name_ptr = 0
|
|
line_num = 0
|
|
_flops = []
|
|
while name_ptr<len(names):
|
|
line = lines[line_num].strip()
|
|
name = names[name_ptr]
|
|
if line.startswith(name):
|
|
flops = float(lines[line_num+1].split(',')[2].strip().split(' ')[0])
|
|
_flops.append(flops)
|
|
name_ptr+=1
|
|
line_num+=1
|
|
|
|
backbone_flops = _flops[:-2]
|
|
neck_flops = _flops[-2]
|
|
head_flops = _flops[-1]
|
|
return all_flops/1e9, backbone_flops, neck_flops, head_flops
|
|
def get_stat(result_dir, group, prefix, idx):
|
|
curr_dir = osp.join(result_dir, group, "%s_%d"%(prefix, idx))
|
|
aps_file = osp.join(curr_dir, 'aps')
|
|
aps = []
|
|
if osp.exists(aps_file):
|
|
with open(aps_file, 'r') as f:
|
|
aps = [float(x) for x in f.readline().strip().split(',')]
|
|
cfg_file = osp.join('configs', group, '%s_%d.py'%(prefix, idx))
|
|
cfg = Config.fromfile(cfg_file)
|
|
all_flops, backbone_flops, neck_flops, head_flops = get_flops(cfg, (3,480,640))
|
|
|
|
return aps, all_flops, backbone_flops, neck_flops, head_flops
|
|
|
|
result_dir = './wouts'
|
|
group = 'scrfdgen2.5g'
|
|
prefix = group
|
|
idx_from = 0
|
|
idx_to = 320
|
|
|
|
outf = open(osp.join(result_dir, "%s.txt"%prefix), 'w')
|
|
|
|
for idx in range(idx_from, idx_to):
|
|
aps, all_flops, backbone_flops, neck_flops, head_flops = get_stat(result_dir, group, prefix, idx)
|
|
backbone_ratio = np.sum(backbone_flops) / all_flops
|
|
neck_ratio = neck_flops / all_flops
|
|
head_ratio = head_flops / all_flops
|
|
print(idx, aps, all_flops, backbone_flops, backbone_ratio, neck_ratio, head_ratio)
|
|
name = "%s_%d"%(prefix, idx)
|
|
data = dict(name=name, backbone_flops=backbone_flops, neck_flops=neck_flops, head_flops=head_flops, all_flops=all_flops, aps=aps)
|
|
data = json.dumps(data)
|
|
outf.write(data)
|
|
outf.write("\n")
|
|
outf.close()
|
|
|