Skip to content

Commit

Permalink
support inference latency when evaluation
Browse files Browse the repository at this point in the history
  • Loading branch information
jihanyang committed Aug 20, 2022
1 parent bd40547 commit ae26b88
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 7 deletions.
20 changes: 17 additions & 3 deletions tools/eval_utils/eval_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,11 @@ def statistics_info(cfg, ret_dict, metric, disp_dict):
'(%d, %d) / %d' % (metric['recall_roi_%s' % str(min_thresh)], metric['recall_rcnn_%s' % str(min_thresh)], metric['gt_num'])


def eval_one_epoch(cfg, model, dataloader, epoch_id, logger, dist_test=False, save_to_file=False, result_dir=None):
def eval_one_epoch(cfg, args, model, dataloader, epoch_id, logger, dist_test=False, result_dir=None):
result_dir.mkdir(parents=True, exist_ok=True)

final_output_dir = result_dir / 'final_result' / 'data'
if save_to_file:
if args.save_to_file:
final_output_dir.mkdir(parents=True, exist_ok=True)

metric = {
Expand All @@ -37,6 +37,10 @@ def eval_one_epoch(cfg, model, dataloader, epoch_id, logger, dist_test=False, sa
class_names = dataset.class_names
det_annos = []

if getattr(args, 'infer_time', False):
start_iter = int(len(dataloader) * 0.1)
infer_time_meter = common_utils.AverageMeter()

logger.info('*************** EPOCH %s EVALUATION *****************' % epoch_id)
if dist_test:
num_gpus = torch.cuda.device_count()
Expand All @@ -53,14 +57,24 @@ def eval_one_epoch(cfg, model, dataloader, epoch_id, logger, dist_test=False, sa
start_time = time.time()
for i, batch_dict in enumerate(dataloader):
load_data_to_gpu(batch_dict)

if getattr(args, 'infer_time', False):
start_time = time.time()

with torch.no_grad():
pred_dicts, ret_dict = model(batch_dict)
disp_dict = {}

if getattr(args, 'infer_time', False):
inference_time = time.time() - start_time
infer_time_meter.update(inference_time * 1000)
# use ms to measure inference time
disp_dict['infer_time'] = f'{infer_time_meter.val:.2f}({infer_time_meter.avg:.2f})'

statistics_info(cfg, ret_dict, metric, disp_dict)
annos = dataset.generate_prediction_dicts(
batch_dict, pred_dicts, class_names,
output_path=final_output_dir if save_to_file else None
output_path=final_output_dir if args.save_to_file else None
)
det_annos += annos
if cfg.LOCAL_RANK == 0:
Expand Down
13 changes: 9 additions & 4 deletions tools/test.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ def parse_config():
parser.add_argument('--eval_all', action='store_true', default=False, help='whether to evaluate all checkpoints')
parser.add_argument('--ckpt_dir', type=str, default=None, help='specify a ckpt directory to be evaluated if needed')
parser.add_argument('--save_to_file', action='store_true', default=False, help='')
parser.add_argument('--infer_time', action='store_true', default=False, help='calculate inference latency')

args = parser.parse_args()

Expand All @@ -60,8 +61,8 @@ def eval_single_ckpt(model, test_loader, args, eval_output_dir, logger, epoch_id

# start evaluation
eval_utils.eval_one_epoch(
cfg, model, test_loader, epoch_id, logger, dist_test=dist_test,
result_dir=eval_output_dir, save_to_file=args.save_to_file
cfg, args, model, test_loader, epoch_id, logger, dist_test=dist_test,
result_dir=eval_output_dir
)


Expand Down Expand Up @@ -118,8 +119,8 @@ def repeat_eval_ckpt(model, test_loader, args, eval_output_dir, logger, ckpt_dir
# start evaluation
cur_result_dir = eval_output_dir / ('epoch_%s' % cur_epoch_id) / cfg.DATA_CONFIG.DATA_SPLIT['test']
tb_dict = eval_utils.eval_one_epoch(
cfg, model, test_loader, cur_epoch_id, logger, dist_test=dist_test,
result_dir=cur_result_dir, save_to_file=args.save_to_file
cfg, model, args, test_loader, cur_epoch_id, logger, dist_test=dist_test,
result_dir=cur_result_dir
)

if cfg.LOCAL_RANK == 0:
Expand All @@ -134,6 +135,10 @@ def repeat_eval_ckpt(model, test_loader, args, eval_output_dir, logger, ckpt_dir

def main():
args, cfg = parse_config()

if args.infer_time:
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

if args.launcher == 'none':
dist_test = False
total_gpus = 1
Expand Down

0 comments on commit ae26b88

Please sign in to comment.