Factor output dir computation out of the testing functions

Reviewed By: rbgirshick

Differential Revision: D7084409

fbshipit-source-id: b1fc559847710e7e4dc35fa8b317c0ba259a4dde

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, gpu_id=0):
+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, gpu_id=0):
         )
     else:
         # Processes entire dataset range by default
-        _boxes, _scores, _ids, rpn_file = generate_rpn_on_range(gpu_id=gpu_id)
+        _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, gpu_id=0):
+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,7 +112,6 @@ def generate_rpn_on_range(ind_range=None, gpu_id=0):
     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))
     )

lib/core/test_engine.py

@@ -81,7 +81,10 @@ def run_inference(ind_range=None, multi_gpu_testing=False, gpu_id=0):
             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)
+            # Note that output dir computation cannot be moved before the
+            # if statement since it depends on the value of cfg.TEST.DATASET
+            output_dir = get_output_dir(training=False)
+            results = parent_func(output_dir, multi_gpu=multi_gpu_testing)
             all_results.update(results)
 
         return all_results
@@ -90,12 +93,12 @@ def run_inference(ind_range=None, multi_gpu_testing=False, gpu_id=0):
         # 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(ind_range=ind_range, gpu_id=gpu_id)
+        output_dir = get_output_dir(training=False)
+        return child_func(output_dir, ind_range=ind_range, gpu_id=gpu_id)
 
 
-def test_net_on_dataset(multi_gpu=False, gpu_id=0):
+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()
@@ -105,7 +108,7 @@ def test_net_on_dataset(multi_gpu=False, gpu_id=0):
             num_images, output_dir
         )
     else:
-        all_boxes, all_segms, all_keyps = test_net(gpu_id=gpu_id)
+        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(
@@ -155,7 +158,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, gpu_id=0):
+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.
     """
@@ -166,7 +169,6 @@ def test_net(ind_range=None, gpu_id=0):
     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
     )