Re-use test_engine multi-gpu testing logic for RetinaNet

Reviewed By: rbgirshick

Differential Revision: D6829206

fbshipit-source-id: 06e8917366a68e4c2e44637d784b019ba2d6fddf

lib/core/test_engine.py

@@ -119,7 +119,9 @@ def test_net(ind_range=None):
         'Use rpn_generate to generate proposals from RPN-only models'
     assert cfg.TEST.DATASET != '', \
         'TEST.DATASET must be set to the dataset name to test'
-
+    if cfg.RETINANET.RETINANET_ON:
+        import core.test_retinanet as test_retinanet
+        return test_retinanet.test_retinanet(ind_range)
     output_dir = get_output_dir(training=False)
     roidb, dataset, start_ind, end_ind, total_num_images = get_roidb_and_dataset(
         ind_range

lib/core/test_retinanet.py

@@ -31,7 +31,6 @@ from caffe2.python import core, workspace
 
 from core.config import cfg, get_output_dir
 from core.rpn_generator import _get_image_blob
-from datasets import task_evaluation
 from datasets.json_dataset import JsonDataset
 from modeling import model_builder
 from modeling.generate_anchors import generate_anchors
@@ -41,9 +40,7 @@ from utils.timer import Timer
 import core.test_engine as test_engine
 import utils.boxes as box_utils
 import utils.c2 as c2_utils
-import utils.env as envu
 import utils.net as nu
-import utils.subprocess as subprocess_utils
 
 logger = logging.getLogger(__name__)
 
@@ -200,7 +197,9 @@ def im_list_detections(model, roidb):
     _t = Timer()
     num_images = len(roidb)
     num_classes = cfg.MODEL.NUM_CLASSES
-    all_boxes, _, _ = test_engine.empty_results(num_classes, num_images)
+    all_boxes, all_segms, all_keyps = test_engine.empty_results(
+        num_classes, num_images
+    )
     # create anchors for each level
     anchors = create_cell_anchors()
     for i, entry in enumerate(roidb):
@@ -213,7 +212,7 @@ def im_list_detections(model, roidb):
         logger.info(
             'im_detections: {:d}/{:d} {:.3f}s'.format(
                 i + 1, num_images, _t.average_time))
-    return all_boxes
+    return all_boxes, all_segms, all_keyps
 
 
 def test_retinanet(ind_range=None):
@@ -238,7 +237,7 @@ def test_retinanet(ind_range=None):
     model_builder.add_inference_inputs(model)
     workspace.CreateNet(model.net)
     # Compute the detections
-    all_boxes = im_list_detections(model, roidb)
+    all_boxes, all_segms, all_keyps = im_list_detections(model, roidb)
     # Save the detections
     cfg_yaml = yaml.dump(cfg)
     if ind_range is not None:
@@ -247,64 +246,12 @@ def test_retinanet(ind_range=None):
         det_name = 'detections.pkl'
     det_file = os.path.join(output_dir, det_name)
     save_object(
-        dict(all_boxes=all_boxes, cfg=cfg_yaml),
-        det_file)
-    logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
-    return all_boxes
-
-
-def multi_gpu_test_retinanet_on_dataset(num_images, output_dir):
-    """
-    If doing multi-gpu testing, we need to divide the data on various gpus and
-    make the subprocess call for each child process that'll run test_retinanet()
-    on its subset data. After all the subprocesses finish, we combine the results
-    and return
-    """
-    # Retrieve the test_net binary path
-    binary_dir = envu.get_runtime_dir()
-    binary_ext = envu.get_py_bin_ext()
-    binary = os.path.join(binary_dir, 'test_net' + binary_ext)
-    assert os.path.exists(binary), 'Binary \'{}\' not found'.format(binary)
-
-    # Run inference in parallel in subprocesses
-    outputs = subprocess_utils.process_in_parallel(
-        'detection', num_images, binary, output_dir)
-
-    # Combine the results from each subprocess
-    all_boxes = [[] for _ in range(cfg.MODEL.NUM_CLASSES)]
-    for det_data in outputs:
-        all_boxes_batch = det_data['all_boxes']
-        for cls_idx in range(1, cfg.MODEL.NUM_CLASSES):
-            all_boxes[cls_idx] += all_boxes_batch[cls_idx]
-
-    # Save the computed detections
-    det_file = os.path.join(output_dir, 'detections.pkl')
-    cfg_yaml = yaml.dump(cfg)
-    save_object(
-        dict(all_boxes=all_boxes, cfg=cfg_yaml),
-        det_file
+        dict(
+            all_boxes=all_boxes,
+            all_segms=all_segms,
+            all_keyps=all_keyps,
+            cfg=cfg_yaml
+        ), det_file
     )
     logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
-
-    return all_boxes
-
-
-def test_retinanet_on_dataset(multi_gpu=False):
-    """
-    Main entry point for testing on a given dataset: whether multi_gpu or not
-    """
-    output_dir = get_output_dir(training=False)
-    dataset = JsonDataset(cfg.TEST.DATASET)
-    test_timer = Timer()
-    test_timer.tic()
-    if multi_gpu:
-        num_images = len(dataset.get_roidb())
-        all_boxes = multi_gpu_test_retinanet_on_dataset(num_images, output_dir)
-    else:
-        all_boxes = test_retinanet()
-    test_timer.toc()
-    logger.info('Total inference time: {:.3f}s'.format(test_timer.average_time))
-    results = task_evaluation.evaluate_all(
-        dataset, all_boxes, None, None, output_dir
-    )
-    return results
+    return all_boxes, all_segms, all_keyps

tools/test_net.py

@@ -38,8 +38,6 @@ from core.config import merge_cfg_from_list
 from core.rpn_generator import generate_rpn_on_dataset
 from core.rpn_generator import generate_rpn_on_range
 from core.test_engine import test_net, test_net_on_dataset
-from core.test_retinanet import test_retinanet
-from core.test_retinanet import test_retinanet_on_dataset
 from datasets import task_evaluation
 import utils.c2
 import utils.logging
@@ -100,12 +98,8 @@ def main(ind_range=None, multi_gpu_testing=False):
     if cfg.MODEL.RPN_ONLY:
         child_func = generate_rpn_on_range
         parent_func = generate_rpn_on_dataset
-    elif cfg.RETINANET.RETINANET_ON:
-        child_func = test_retinanet
-        parent_func = test_retinanet_on_dataset
     else:
         # Generic case that handles all network types other than RPN-only nets
-        # and RetinaNet
         child_func = test_net
         parent_func = test_net_on_dataset