Do not mutate TEST.WEIGHTS

Reviewed By: ir413

Differential Revision: D7148432

fbshipit-source-id: e23f767fe0241496b0ee0fb73c247747250eb34e

lib/core/rpn_generator.py

@@ -54,7 +54,12 @@ logger = logging.getLogger(__name__)
 
 
 def generate_rpn_on_dataset(
-    dataset_name, _proposal_file_ignored, output_dir, multi_gpu=False, gpu_id=0
+    weights_file,
+    dataset_name,
+    _proposal_file_ignored,
+    output_dir,
+    multi_gpu=False,
+    gpu_id=0
 ):
     """Run inference on a dataset."""
     dataset = JsonDataset(dataset_name)
@@ -63,12 +68,17 @@ def generate_rpn_on_dataset(
     if multi_gpu:
         num_images = len(dataset.get_roidb())
         _boxes, _scores, _ids, rpn_file = multi_gpu_generate_rpn_on_dataset(
-            dataset_name, _proposal_file_ignored, num_images, output_dir
+            weights_file, dataset_name, _proposal_file_ignored, num_images,
+            output_dir
         )
     else:
         # Processes entire dataset range by default
         _boxes, _scores, _ids, rpn_file = generate_rpn_on_range(
-            dataset_name, _proposal_file_ignored, output_dir, gpu_id=gpu_id
+            weights_file,
+            dataset_name,
+            _proposal_file_ignored,
+            output_dir,
+            gpu_id=gpu_id
         )
     test_timer.toc()
     logger.info('Total inference time: {:.3f}s'.format(test_timer.average_time))
@@ -76,7 +86,7 @@ def generate_rpn_on_dataset(
 
 
 def multi_gpu_generate_rpn_on_dataset(
-    dataset_name, _proposal_file_ignored, num_images, output_dir
+    weights_file, dataset_name, _proposal_file_ignored, num_images, output_dir
 ):
     """Multi-gpu inference on a dataset."""
     # Retrieve the test_net binary path
@@ -87,6 +97,7 @@ def multi_gpu_generate_rpn_on_dataset(
 
     # Pass the target dataset via the command line
     opts = ['TEST.DATASETS', '("{}",)'.format(dataset_name)]
+    opts += ['TEST.WEIGHTS', weights_file]
 
     # Run inference in parallel in subprocesses
     outputs = subprocess_utils.process_in_parallel(
@@ -109,13 +120,16 @@ def multi_gpu_generate_rpn_on_dataset(
 
 
 def generate_rpn_on_range(
-    dataset_name, _proposal_file_ignored, output_dir, ind_range=None, gpu_id=0
+    weights_file,
+    dataset_name,
+    _proposal_file_ignored,
+    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.
     """
-    assert cfg.TEST.WEIGHTS != '', \
-        'TEST.WEIGHTS must be set to the model file to test'
     assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
 
     roidb, start_ind, end_ind, total_num_images = get_roidb(
@@ -127,7 +141,7 @@ def generate_rpn_on_range(
 
     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, weights_file, gpu_id=gpu_id,
     )
     model_builder.add_inference_inputs(model)
     workspace.CreateNet(model.net)

lib/core/test_engine.py

@@ -82,7 +82,9 @@ def get_inference_dataset(index, is_parent=True):
     return dataset_name, proposal_file
 
 
-def run_inference(ind_range=None, multi_gpu_testing=False, gpu_id=0):
+def run_inference(
+    weights_file, ind_range=None, multi_gpu_testing=False, gpu_id=0
+):
     parent_func, child_func = get_eval_functions()
 
     is_parent = ind_range is None
@@ -96,6 +98,7 @@ def run_inference(ind_range=None, multi_gpu_testing=False, gpu_id=0):
             dataset_name, proposal_file = get_inference_dataset(i)
             output_dir = get_output_dir(dataset_name, training=False)
             results = parent_func(
+                weights_file,
                 dataset_name,
                 proposal_file,
                 output_dir,
@@ -111,6 +114,7 @@ def run_inference(ind_range=None, multi_gpu_testing=False, gpu_id=0):
         dataset_name, proposal_file = get_inference_dataset(0, is_parent=False)
         output_dir = get_output_dir(dataset_name, training=False)
         return child_func(
+            weights_file,
             dataset_name,
             proposal_file,
             output_dir,
@@ -120,7 +124,12 @@ def run_inference(ind_range=None, multi_gpu_testing=False, gpu_id=0):
 
 
 def test_net_on_dataset(
-    dataset_name, proposal_file, output_dir, multi_gpu=False, gpu_id=0
+    weights_file,
+    dataset_name,
+    proposal_file,
+    output_dir,
+    multi_gpu=False,
+    gpu_id=0
 ):
     """Run inference on a dataset."""
     dataset = JsonDataset(dataset_name)
@@ -129,11 +138,11 @@ def test_net_on_dataset(
     if multi_gpu:
         num_images = len(dataset.get_roidb())
         all_boxes, all_segms, all_keyps = multi_gpu_test_net_on_dataset(
-            dataset_name, proposal_file, num_images, output_dir
+            weights_file, dataset_name, proposal_file, num_images, output_dir
         )
     else:
         all_boxes, all_segms, all_keyps = test_net(
-            dataset_name, proposal_file, output_dir, gpu_id=gpu_id
+            weights_file, dataset_name, proposal_file, output_dir, gpu_id=gpu_id
         )
     test_timer.toc()
     logger.info('Total inference time: {:.3f}s'.format(test_timer.average_time))
@@ -144,7 +153,7 @@ def test_net_on_dataset(
 
 
 def multi_gpu_test_net_on_dataset(
-    dataset_name, proposal_file, num_images, output_dir
+    weights_file, dataset_name, proposal_file, num_images, output_dir
 ):
     """Multi-gpu inference on a dataset."""
     binary_dir = envu.get_runtime_dir()
@@ -154,6 +163,7 @@ def multi_gpu_test_net_on_dataset(
 
     # Pass the target dataset and proposal file (if any) via the command line
     opts = ['TEST.DATASETS', '("{}",)'.format(dataset_name)]
+    opts += ['TEST.WEIGHTS', weights_file]
     if proposal_file:
         opts += ['TEST.PROPOSAL_FILES', '("{}",)'.format(proposal_file)]
 
@@ -191,19 +201,24 @@ def multi_gpu_test_net_on_dataset(
     return all_boxes, all_segms, all_keyps
 
 
-def test_net(dataset_name, proposal_file, output_dir, ind_range=None, gpu_id=0):
+def test_net(
+    weights_file,
+    dataset_name,
+    proposal_file,
+    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.
     """
-    assert cfg.TEST.WEIGHTS != '', \
-        'TEST.WEIGHTS must be set to the model file to test'
     assert not cfg.MODEL.RPN_ONLY, \
         'Use rpn_generate to generate proposals from RPN-only models'
 
     roidb, dataset, start_ind, end_ind, total_num_images = get_roidb_and_dataset(
         dataset_name, proposal_file, ind_range
     )
-    model = initialize_model_from_cfg(gpu_id=gpu_id)
+    model = initialize_model_from_cfg(weights_file, 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)
@@ -292,13 +307,13 @@ def test_net(dataset_name, proposal_file, output_dir, ind_range=None, gpu_id=0):
     return all_boxes, all_segms, all_keyps
 
 
-def initialize_model_from_cfg(gpu_id=0):
+def initialize_model_from_cfg(weights_file, 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, gpu_id=gpu_id)
     net_utils.initialize_gpu_from_weights_file(
-        model, cfg.TEST.WEIGHTS, gpu_id=gpu_id,
+        model, weights_file, gpu_id=gpu_id,
     )
     model_builder.add_inference_inputs(model)
     workspace.CreateNet(model.net)

tools/infer.py

@@ -98,14 +98,13 @@ def parse_args():
 
 def get_rpn_box_proposals(im, args):
     merge_cfg_from_file(args.rpn_cfg)
-    cfg.TEST.WEIGHTS = args.rpn_pkl
     cfg.NUM_GPUS = 1
     cfg.MODEL.RPN_ONLY = True
     cfg.TEST.RPN_PRE_NMS_TOP_N = 10000
     cfg.TEST.RPN_POST_NMS_TOP_N = 2000
     assert_and_infer_cfg()
 
-    model = model_engine.initialize_model_from_cfg()
+    model = model_engine.initialize_model_from_cfg(args.rpn_pkl)
     with c2_utils.NamedCudaScope(0):
         boxes, scores = rpn_engine.im_proposals(model, im)
     return boxes, scores
@@ -129,10 +128,12 @@ def main(args):
         merge_cfg_from_cfg(cfg_orig)
         merge_cfg_from_file(yml)
         if len(pkl) > 0:
-            cfg.TEST.WEIGHTS = pkl
+            weights_file = pkl
+        else:
+            weights_file = cfg.TEST.WEIGHTS
         cfg.NUM_GPUS = 1
         assert_and_infer_cfg()
-        model = model_engine.initialize_model_from_cfg()
+        model = model_engine.initialize_model_from_cfg(weights_file)
         with c2_utils.NamedCudaScope(0):
             cls_boxes_, cls_segms_, cls_keyps_ = \
                 model_engine.im_detect_all(model, im, proposal_boxes)

tools/infer_simple.py

@@ -93,10 +93,9 @@ def parse_args():
 def main(args):
     logger = logging.getLogger(__name__)
     merge_cfg_from_file(args.cfg)
-    cfg.TEST.WEIGHTS = args.weights
     cfg.NUM_GPUS = 1
     assert_and_infer_cfg()
-    model = infer_engine.initialize_model_from_cfg()
+    model = infer_engine.initialize_model_from_cfg(args.weights)
     dummy_coco_dataset = dummy_datasets.get_coco_dataset()
 
     if os.path.isdir(args.im_or_folder):

tools/test_net.py

@@ -91,9 +91,11 @@ def parse_args():
     return parser.parse_args()
 
 
-def main(ind_range=None, multi_gpu_testing=False):
+def main(weights_file, ind_range=None, multi_gpu_testing=False):
     all_results = run_inference(
-        ind_range=ind_range, multi_gpu_testing=multi_gpu_testing
+        weights_file,
+        ind_range=ind_range,
+        multi_gpu_testing=multi_gpu_testing,
     )
     if not ind_range:
         task_evaluation.check_expected_results(
@@ -122,4 +124,8 @@ if __name__ == '__main__':
         logger.info('Waiting for \'{}\' to exist...'.format(cfg.TEST.WEIGHTS))
         time.sleep(10)
 
-    main(ind_range=args.range, multi_gpu_testing=args.multi_gpu_testing)
+    main(
+        cfg.TEST.WEIGHTS,
+        ind_range=args.range,
+        multi_gpu_testing=args.multi_gpu_testing
+    )

tools/train_net.py

@@ -270,12 +270,10 @@ def dump_proto_files(model, output_dir):
 
 def test_model(model_file, multi_gpu_testing, opts=None):
     """Test a model."""
-    # All arguments to inference functions are passed via cfg
-    cfg.TEST.WEIGHTS = model_file
     # Clear memory before inference
     workspace.ResetWorkspace()
     # Run inference
-    test_net.main(multi_gpu_testing=multi_gpu_testing)
+    test_net.main(model_file, multi_gpu_testing=multi_gpu_testing)
 
 
 if __name__ == '__main__':