Allow for multi-threaded evaluation code

Summary:
Under certain settings, it is desirable to use multi-threaded
evaluation even though it is inefficient due to Python's GIL. This set of
changes allows for it.

Reviewed By: rbgirshick

Differential Revision: D6827125

fbshipit-source-id: 845647bd09198c35bb6d06c1a453a817d0dcf784

lib/core/rpn_generator.py

@@ -39,7 +39,6 @@ from caffe2.python import core
 from caffe2.python import workspace
 
 from core.config import cfg
-from core.config import get_output_dir
 from datasets import task_evaluation
 from datasets.json_dataset import JsonDataset
 from modeling import model_builder
@@ -54,9 +53,8 @@ import utils.subprocess as subprocess_utils
 logger = logging.getLogger(__name__)
 
 
-def generate_rpn_on_dataset(multi_gpu=False):
+def generate_rpn_on_dataset(output_dir, multi_gpu=False, gpu_id=0):
     """Run inference on a dataset."""
-    output_dir = get_output_dir(training=False)
     dataset = JsonDataset(cfg.TEST.DATASET)
     test_timer = Timer()
     test_timer.tic()
@@ -67,7 +65,9 @@ def generate_rpn_on_dataset(multi_gpu=False):
         )
     else:
         # Processes entire dataset range by default
-        _boxes, _scores, _ids, rpn_file = generate_rpn_on_range()
+        _boxes, _scores, _ids, rpn_file = generate_rpn_on_range(
+            output_dir, gpu_id=gpu_id
+        )
     test_timer.toc()
     logger.info('Total inference time: {:.3f}s'.format(test_timer.average_time))
     return evaluate_proposal_file(dataset, rpn_file, output_dir)
@@ -101,7 +101,7 @@ def multi_gpu_generate_rpn_on_dataset(num_images, output_dir):
     return boxes, scores, ids, rpn_file
 
 
-def generate_rpn_on_range(ind_range=None):
+def generate_rpn_on_range(output_dir, ind_range=None, gpu_id=0):
     """Run inference on all images in a dataset or over an index range of images
     in a dataset using a single GPU.
     """
@@ -112,13 +112,14 @@ def generate_rpn_on_range(ind_range=None):
     assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
 
     roidb, start_ind, end_ind, total_num_images = get_roidb(ind_range)
-    output_dir = get_output_dir(training=False)
     logger.info(
         'Output will be saved to: {:s}'.format(os.path.abspath(output_dir))
     )
 
-    model = model_builder.create(cfg.MODEL.TYPE, train=False)
-    nu.initialize_from_weights_file(model, cfg.TEST.WEIGHTS)
+    model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
+    nu.initialize_gpu_from_weights_file(
+        model, cfg.TEST.WEIGHTS, gpu_id=gpu_id,
+    )
     model_builder.add_inference_inputs(model)
     workspace.CreateNet(model.net)
 
@@ -127,7 +128,8 @@ def generate_rpn_on_range(ind_range=None):
         roidb,
         start_ind=start_ind,
         end_ind=end_ind,
-        total_num_images=total_num_images
+        total_num_images=total_num_images,
+        gpu_id=gpu_id,
     )
 
     cfg_yaml = yaml.dump(cfg)
@@ -144,7 +146,8 @@ def generate_rpn_on_range(ind_range=None):
 
 
 def generate_proposals_on_roidb(
-    model, roidb, start_ind=None, end_ind=None, total_num_images=None
+    model, roidb, start_ind=None, end_ind=None, total_num_images=None,
+    gpu_id=0,
 ):
     """Generate RPN proposals on all images in an imdb."""
     _t = Timer()
@@ -159,7 +162,7 @@ def generate_proposals_on_roidb(
     for i in range(num_images):
         roidb_ids[i] = roidb[i]['id']
         im = cv2.imread(roidb[i]['image'])
-        with c2_utils.NamedCudaScope(0):
+        with c2_utils.NamedCudaScope(gpu_id):
             _t.tic()
             roidb_boxes[i], roidb_scores[i] = im_proposals(model, im)
             _t.toc()

lib/core/test_engine.py

@@ -31,7 +31,8 @@ import yaml
 from caffe2.python import workspace
 
 from core.config import cfg
-from core.config import get_output_dir
+from core.rpn_generator import generate_rpn_on_dataset
+from core.rpn_generator import generate_rpn_on_range
 from core.test import im_detect_all
 from datasets import task_evaluation
 from datasets.json_dataset import JsonDataset
@@ -47,9 +48,52 @@ import utils.vis as vis_utils
 logger = logging.getLogger(__name__)
 
 
-def test_net_on_dataset(multi_gpu=False):
+def get_eval_functions():
+    # Determine which parent or child function should handle inference
+    if cfg.MODEL.RPN_ONLY:
+        child_func = generate_rpn_on_range
+        parent_func = generate_rpn_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
+
+    return parent_func, child_func
+
+
+def run_inference(output_dir, ind_range=None, multi_gpu_testing=False, gpu_id=0):
+    parent_func, child_func = get_eval_functions()
+
+    is_parent = ind_range is None
+    if is_parent:
+        # Parent case:
+        # In this case we're either running inference on the entire dataset in a
+        # single process or (if multi_gpu_testing is True) using this process to
+        # launch subprocesses that each run inference on a range of the dataset
+        if len(cfg.TEST.DATASETS) == 0:
+            cfg.TEST.DATASETS = (cfg.TEST.DATASET, )
+            cfg.TEST.PROPOSAL_FILES = (cfg.TEST.PROPOSAL_FILE, )
+
+        all_results = {}
+        for i in range(len(cfg.TEST.DATASETS)):
+            cfg.TEST.DATASET = cfg.TEST.DATASETS[i]
+            if cfg.TEST.PRECOMPUTED_PROPOSALS:
+                cfg.TEST.PROPOSAL_FILE = cfg.TEST.PROPOSAL_FILES[i]
+            results = parent_func(output_dir, multi_gpu=multi_gpu_testing)
+            all_results.update(results)
+
+        return all_results
+    else:
+        # Subprocess child case:
+        # In this case test_net was called via subprocess.Popen to execute on a
+        # range of inputs on a single dataset (i.e., use cfg.TEST.DATASET and
+        # don't loop over cfg.TEST.DATASETS)
+        return child_func(output_dir, ind_range=ind_range, gpu_id=gpu_id)
+
+
+def test_net_on_dataset(output_dir, multi_gpu=False, gpu_id=0):
     """Run inference on a dataset."""
-    output_dir = get_output_dir(training=False)
     dataset = JsonDataset(cfg.TEST.DATASET)
     test_timer = Timer()
     test_timer.tic()
@@ -59,7 +103,7 @@ def test_net_on_dataset(multi_gpu=False):
             num_images, output_dir
         )
     else:
-        all_boxes, all_segms, all_keyps = test_net()
+        all_boxes, all_segms, all_keyps = test_net(output_dir, gpu_id=gpu_id)
     test_timer.toc()
     logger.info('Total inference time: {:.3f}s'.format(test_timer.average_time))
     results = task_evaluation.evaluate_all(
@@ -109,7 +153,7 @@ def multi_gpu_test_net_on_dataset(num_images, output_dir):
     return all_boxes, all_segms, all_keyps
 
 
-def test_net(ind_range=None):
+def test_net(output_dir, ind_range=None, gpu_id=0):
     """Run inference on all images in a dataset or over an index range of images
     in a dataset using a single GPU.
     """
@@ -119,11 +163,11 @@ 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'
-    output_dir = get_output_dir(training=False)
+
     roidb, dataset, start_ind, end_ind, total_num_images = get_roidb_and_dataset(
         ind_range
     )
-    model = initialize_model_from_cfg()
+    model = initialize_model_from_cfg(gpu_id=gpu_id)
     num_images = len(roidb)
     num_classes = cfg.MODEL.NUM_CLASSES
     all_boxes, all_segms, all_keyps = empty_results(num_classes, num_images)
@@ -144,7 +188,7 @@ def test_net(ind_range=None):
             box_proposals = None
 
         im = cv2.imread(entry['image'])
-        with c2_utils.NamedCudaScope(0):
+        with c2_utils.NamedCudaScope(gpu_id):
             cls_boxes_i, cls_segms_i, cls_keyps_i = im_detect_all(
                 model, im, box_proposals, timers
             )
@@ -212,13 +256,13 @@ def test_net(ind_range=None):
     return all_boxes, all_segms, all_keyps
 
 
-def initialize_model_from_cfg():
+def initialize_model_from_cfg(gpu_id=0):
     """Initialize a model from the global cfg. Loads test-time weights and
     creates the networks in the Caffe2 workspace.
     """
-    model = model_builder.create(cfg.MODEL.TYPE, train=False)
-    net_utils.initialize_from_weights_file(
-        model, cfg.TEST.WEIGHTS, broadcast=False
+    model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
+    net_utils.initialize_gpu_from_weights_file(
+        model, cfg.TEST.WEIGHTS, gpu_id=gpu_id,
     )
     model_builder.add_inference_inputs(model)
     workspace.CreateNet(model.net)

tools/test_net.py

@@ -33,11 +33,10 @@ from caffe2.python import workspace
 
 from core.config import assert_and_infer_cfg
 from core.config import cfg
+from core.config import get_output_dir
 from core.config import merge_cfg_from_file
 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_engine import run_inference
 from datasets import task_evaluation
 import utils.c2
 import utils.logging
@@ -94,46 +93,17 @@ def parse_args():
 
 
 def main(ind_range=None, multi_gpu_testing=False):
-    # Determine which parent or child function should handle inference
-    if cfg.MODEL.RPN_ONLY:
-        child_func = generate_rpn_on_range
-        parent_func = generate_rpn_on_dataset
-    else:
-        # Generic case that handles all network types other than RPN-only nets
-        child_func = test_net
-        parent_func = test_net_on_dataset
-
-    is_parent = ind_range is None
-
-    if is_parent:
-        # Parent case:
-        # In this case we're either running inference on the entire dataset in a
-        # single process or (if multi_gpu_testing is True) using this process to
-        # launch subprocesses that each run inference on a range of the dataset
-        if len(cfg.TEST.DATASETS) == 0:
-            cfg.TEST.DATASETS = (cfg.TEST.DATASET, )
-            cfg.TEST.PROPOSAL_FILES = (cfg.TEST.PROPOSAL_FILE, )
-
-        all_results = {}
-        for i in range(len(cfg.TEST.DATASETS)):
-            cfg.TEST.DATASET = cfg.TEST.DATASETS[i]
-            if cfg.TEST.PRECOMPUTED_PROPOSALS:
-                cfg.TEST.PROPOSAL_FILE = cfg.TEST.PROPOSAL_FILES[i]
-            results = parent_func(multi_gpu=multi_gpu_testing)
-            all_results.update(results)
-
+    output_dir = get_output_dir(training=False)
+    all_results = run_inference(
+        output_dir, ind_range=ind_range, multi_gpu_testing=multi_gpu_testing
+    )
+    if not ind_range:
         task_evaluation.check_expected_results(
             all_results,
             atol=cfg.EXPECTED_RESULTS_ATOL,
             rtol=cfg.EXPECTED_RESULTS_RTOL
         )
         task_evaluation.log_copy_paste_friendly_results(all_results)
-    else:
-        # Subprocess child case:
-        # In this case test_net was called via subprocess.Popen to execute on a
-        # range of inputs on a single dataset (i.e., use cfg.TEST.DATASET and
-        # don't loop over cfg.TEST.DATASETS)
-        child_func(ind_range=ind_range)
 
 
 if __name__ == '__main__':