support inference latency when evaluation

tools/eval_utils/eval_utils.py

@@ -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 = {
@@ -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()
@@ -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:

tools/test.py

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