def create_fake_detection_data_loader(height, width, is_train): with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir: runner = create_runner("d2go.runner.GeneralizedRCNNRunner") cfg = runner.get_default_cfg() cfg.DATASETS.TRAIN = ["default_dataset_train"] cfg.DATASETS.TEST = ["default_dataset_test"] with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir: image_dir = os.path.join(dataset_dir, "images") os.makedirs(image_dir) image_generator = LocalImageGenerator(image_dir, width=width, height=height) if is_train: with register_toy_dataset("default_dataset_train", image_generator, num_images=3): train_loader = runner.build_detection_train_loader(cfg) yield train_loader else: with register_toy_dataset("default_dataset_test", image_generator, num_images=3): test_loader = runner.build_detection_test_loader( cfg, dataset_name="default_dataset_test") yield test_loader
def test_script_only_model(self): def _validate(predictor): outputs = predictor( [torch.tensor(1), torch.tensor(2), torch.tensor(3)]) self.assertEqual(len(outputs), 3) self.assertEqual( outputs, [torch.tensor(0), torch.tensor(2), torch.tensor(6)]) # Method 1: explicitly set jit_mode to "trace" with make_temp_directory("test_test_script_only_model") as tmp_dir: model = ScriptingOnlyModel() predictor = self._export_simple_model("explicit", model, None, tmp_dir, predictor_type="torchscript") _validate(predictor) # Method 2: using torchscript@scripting as predictor type with make_temp_directory("test_test_script_only_model") as tmp_dir: model = ScriptingOnlyModel() predictor = self._export_simple_model( "implicit", model, None, tmp_dir, predictor_type="torchscript@scripting") _validate(predictor)
def _register_toy_dataset( dataset_name, image_generator, num_images, num_classes=-1, num_keypoints=0 ): json_dataset, meta_data = create_toy_dataset( image_generator, num_images=num_images, num_classes=num_classes, num_keypoints=num_keypoints, ) with make_temp_directory("detectron2go_tmp_dataset") as tmp_dir: json_file = os.path.join(tmp_dir, "{}.json".format(dataset_name)) with open(json_file, "w") as f: json.dump(json_dataset, f) split_dict = { IM_DIR: image_generator.get_image_dir(), ANN_FN: json_file, "meta_data": meta_data, } register_dataset_split(dataset_name, split_dict) try: yield finally: DatasetCatalog.remove(dataset_name) MetadataCatalog.remove(dataset_name)
def test_export_and_load( cls, model, input_args, export_method, export_kwargs, output_checker ): """ Illustrate the life-cycle of export and load, used for testing. """ with make_temp_directory("test_export_and_load") as save_path: # run the orginal model assert isinstance(model, nn.Module), model assert isinstance(input_args, (list, tuple)), input_args original_output = model(*input_args) # export the model model.eval() # TODO: decide where eval() should be called load_kwargs = cls.export( model, input_args, save_path, export_method, **export_kwargs ) # sanity check for load_kwargs assert isinstance(load_kwargs, dict), load_kwargs assert json.dumps(load_kwargs), load_kwargs # loaded model back loaded_model = cls.load(save_path, **load_kwargs) # run the loaded model assert isinstance(loaded_model, nn.Module), loaded_model new_output = loaded_model(*input_args) # compare outputs output_checker(new_output, original_output)
def test_coco_conversions(self): test_data_0 = { "info": {}, "imgs": { "img_1": { "file_name": "0.jpg", "width": 600, "height": 600, "id": "img_1", } }, "anns": {0: {"id": 0, "image_id": "img_1", "bbox": [30, 30, 60, 20]}}, "imgToAnns": {"img_1": [0]}, "cats": {}, } test_data_1 = copy.deepcopy(test_data_0) test_data_1["imgs"][123] = test_data_1["imgs"].pop("img_1") test_data_1["imgs"][123]["id"] = 123 test_data_1["anns"][0]["image_id"] = 123 test_data_1["imgToAnns"][123] = test_data_1["imgToAnns"].pop("img_1") for test_data, exp_output in [(test_data_0, [0, 0]), (test_data_1, [123, 123])]: with make_temp_directory("detectron2go_tmp_dataset") as tmp_dir: src_json = os.path.join(tmp_dir, "source.json") out_json = os.path.join(tmp_dir, "output.json") with open(src_json, "w") as h_in: json.dump(test_data, h_in) out_json = extended_coco.convert_coco_text_to_coco_detection_json( src_json, out_json ) self.assertEqual(out_json["images"][0]["id"], exp_output[0]) self.assertEqual(out_json["annotations"][0]["image_id"], exp_output[1])
def test_build_model(self): cfg = self._get_default_cfg() cfg.INPUT.MIN_SIZE_TRAIN = (60,) cfg.MODEL.KMEANS_ANCHORS.KMEANS_ANCHORS_ON = True cfg.MODEL.KMEANS_ANCHORS.NUM_CLUSTERS = 3 cfg.MODEL.KMEANS_ANCHORS.NUM_TRAINING_IMG = 5 cfg.MODEL.KMEANS_ANCHORS.DATASETS = ("toy_dataset",) cfg.MODEL.DEVICE = "cpu" cfg.MODEL.ANCHOR_GENERATOR.NAME = "KMeansAnchorGenerator" with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir: image_dir = os.path.join(dataset_dir, "images") os.makedirs(image_dir) image_generator = LocalImageGenerator(image_dir, width=80, height=60) with register_toy_dataset( "toy_dataset", image_generator, num_images=cfg.MODEL.KMEANS_ANCHORS.NUM_TRAINING_IMG, ): model = self.runner.build_model(cfg) trainer = SimpleTrainer(model, data_loader=[], optimizer=None) trainer_hooks = [compute_kmeans_anchors_hook(self.runner, cfg)] trainer.register_hooks(trainer_hooks) trainer.before_train() anchor_generator = model.proposal_generator.anchor_generator cell_anchors = [x for x in anchor_generator.cell_anchors] gt_anchors = np.array( [ [-20, -15, 20, 15] # toy_dataset's bbox is half size of image for _ in range(cfg.MODEL.KMEANS_ANCHORS.NUM_CLUSTERS) ] ) np.testing.assert_allclose(cell_anchors[0], gt_anchors)
def test_coco_injection(self): with make_temp_directory("detectron2go_tmp_dataset") as tmp_dir: image_dir, json_file = create_test_images_and_dataset_json(tmp_dir) runner = Detectron2GoRunner() cfg = runner.get_default_cfg() cfg.merge_from_list( [ str(x) for x in [ "D2GO_DATA.DATASETS.COCO_INJECTION.NAMES", ["inj_ds1", "inj_ds2"], "D2GO_DATA.DATASETS.COCO_INJECTION.IM_DIRS", [image_dir, "/mnt/fair"], "D2GO_DATA.DATASETS.COCO_INJECTION.JSON_FILES", [json_file, "inj_ds2"], ] ] ) runner.register(cfg) inj_ds1 = DatasetCatalog.get("inj_ds1") self.assertEqual(len(inj_ds1), 10) for dic in inj_ds1: self.assertEqual(dic["width"], 80) self.assertEqual(dic["height"], 60)
def test_sub_dataset(self): with make_temp_directory("detectron2go_tmp_dataset") as tmp_dir: image_dir, json_file = create_test_images_and_dataset_json(tmp_dir) runner = Detectron2GoRunner() cfg = runner.get_default_cfg() cfg.merge_from_list( [ str(x) for x in [ "D2GO_DATA.DATASETS.COCO_INJECTION.NAMES", ["inj_ds"], "D2GO_DATA.DATASETS.COCO_INJECTION.IM_DIRS", [image_dir], "D2GO_DATA.DATASETS.COCO_INJECTION.JSON_FILES", [json_file], "DATASETS.TEST", ("inj_ds",), "D2GO_DATA.TEST.MAX_IMAGES", 1, ] ] ) runner.register(cfg) with maybe_subsample_n_images(cfg) as new_cfg: test_loader = runner.build_detection_test_loader( new_cfg, new_cfg.DATASETS.TEST[0] ) self.assertEqual(len(test_loader), 1)
def test_default_dataset(self): runner = create_runner("d2go.runner.GeneralizedRCNNRunner") cfg = runner.get_default_cfg() cfg.DATASETS.TRAIN = ["default_dataset_train"] cfg.DATASETS.TEST = ["default_dataset_test"] with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir: image_dir = os.path.join(dataset_dir, "images") os.makedirs(image_dir) image_generator = LocalImageGenerator(image_dir, width=80, height=60) with register_toy_dataset("default_dataset_train", image_generator, num_images=3): train_loader = runner.build_detection_train_loader(cfg) for i, data in enumerate(train_loader): self.assertIsNotNone(data) # for training loader, it has infinite length if i == 6: break with register_toy_dataset("default_dataset_test", image_generator, num_images=3): test_loader = runner.build_detection_test_loader( cfg, dataset_name="default_dataset_test") all_data = [] for data in test_loader: all_data.append(data) self.assertEqual(len(all_data), 3)
def _test_export(self, predictor_type, compare_match=True): size_divisibility = max(self.test_model.backbone.size_divisibility, 10) h, w = size_divisibility, size_divisibility * 2 with create_fake_detection_data_loader(h, w, is_train=False) as data_loader: inputs = next(iter(data_loader)) with make_temp_directory( "test_export_{}".format(predictor_type) ) as tmp_dir: # TODO: the export may change model it self, need to fix this model_to_export = copy.deepcopy(self.test_model) predictor_path = convert_and_export_predictor( self.cfg, model_to_export, predictor_type, tmp_dir, data_loader ) predictor = create_predictor(predictor_path) predicotr_outputs = predictor(inputs) _validate_outputs(inputs, predicotr_outputs) if compare_match: with torch.no_grad(): pytorch_outputs = self.test_model(inputs) assert_instances_allclose( predicotr_outputs[0]["instances"], pytorch_outputs[0]["instances"], )
def test_fetch_checkpoints_local(self): with make_temp_directory("test") as output_dir: output_dir = Path(output_dir) for i in range(5): create_file(output_dir / f"model_{i}.pth") create_file(output_dir / "model_final.pth") checkpoints = list(fetch_checkpoints_till_final(output_dir)) assert len(checkpoints) == 6
def test_import_from_file(self): with make_temp_directory("test_import_from_file") as dir: filename = os.path.join(dir, "my_module.py") with open(filename, "w") as f: f.write("ANSWER = 42\n") my_module = import_file("my_module", filename) self.assertEqual(my_module.ANSWER, 42)
def trace_and_save_torchscript( model: nn.Module, inputs: Tuple[torch.Tensor], output_path: str, mobile_optimization: Optional[MobileOptimizationConfig] = None, _extra_files: Optional[Dict[str, bytes]] = None, ): logger.info("Tracing and saving TorchScript to {} ...".format(output_path)) PathManager.mkdirs(output_path) if _extra_files is None: _extra_files = {} # TODO: patch_builtin_len depends on D2, we should either copy the function or # dynamically registering the D2's version. from detectron2.export.torchscript_patch import patch_builtin_len with torch.no_grad(), patch_builtin_len(): script_model = torch.jit.trace(model, inputs) with make_temp_directory("trace_and_save_torchscript") as tmp_dir: @contextlib.contextmanager def _synced_local_file(rel_path): remote_file = os.path.join(output_path, rel_path) local_file = os.path.join(tmp_dir, rel_path) yield local_file PathManager.copy_from_local(local_file, remote_file, overwrite=True) with _synced_local_file("model.jit") as model_file: torch.jit.save(script_model, model_file, _extra_files=_extra_files) with _synced_local_file("data.pth") as data_file: torch.save(inputs, data_file) if mobile_optimization is not None: logger.info("Applying optimize_for_mobile ...") liteopt_model = optimize_for_mobile( script_model, optimization_blocklist=mobile_optimization. optimization_blocklist, preserved_methods=mobile_optimization.preserved_methods, backend=mobile_optimization.backend, ) with _synced_local_file("mobile_optimized.ptl") as lite_path: liteopt_model._save_for_lite_interpreter(lite_path) # liteopt_model(*inputs) # sanity check op_names = torch.jit.export_opnames(liteopt_model) logger.info("Operator names from lite interpreter:\n{}".format( "\n".join(op_names))) logger.info("Applying augment_model_with_bundled_inputs ...") augment_model_with_bundled_inputs(liteopt_model, [inputs]) liteopt_model.run_on_bundled_input(0) # sanity check with _synced_local_file( "mobile_optimized_bundled.ptl") as lite_path: liteopt_model._save_for_lite_interpreter(lite_path)
def test_model_info(self): with make_temp_directory("test_model_info") as tmp_dir: _save_test_model(tmp_dir) model_info = ModelInfo(path=tmp_dir, type="torchscript") # NOTE: decide if load_model is a public API or class method of ModelInfo from mobile_cv.predictor.model_wrappers import load_model model = load_model(model_info, model_root="") self.assertEqual(torch.tensor(2), model(torch.tensor(1)))
def test_simple_two_part_model(self): with make_temp_directory("test_simple_two_part_model") as tmp_dir: model = TwoPartSimpleModel() predictor = self._export_simple_model(None, model, torch.tensor(1), tmp_dir, predictor_type="torchscript") x = torch.tensor(42) self.assertEqual(predictor(x), model(x))
def test_fetch_lightning_checkpoints_local(self): with make_temp_directory("test") as output_dir: output_dir = Path(output_dir) ext = ModelCheckpoint.FILE_EXTENSION for i in range(5): create_file(output_dir / f"step={i}{ext}") create_file(output_dir / f"model_final{ext}") create_file(output_dir / f"{ModelCheckpoint.CHECKPOINT_NAME_LAST}{ext}") checkpoints = list(fetch_checkpoints_till_final(output_dir)) self.assertEqual(len(checkpoints), 6)
def test_export_torchvision_format(self): runner = GeneralizedRCNNRunner() cfg = runner.get_default_cfg() cfg.merge_from_file("detectron2go://mask_rcnn_fbnetv3a_dsmask_C4.yaml") cfg.merge_from_list(get_quick_test_config_opts()) cfg.merge_from_list(["MODEL.DEVICE", "cpu"]) pytorch_model = runner.build_model(cfg, eval_only=True) from typing import Dict, List class Wrapper(torch.nn.Module): def __init__(self, model): super().__init__() self.model = model def forward(self, inputs: List[torch.Tensor]): x = inputs[0].unsqueeze(0) * 255 scale = 320.0 / min(x.shape[-2], x.shape[-1]) x = torch.nn.functional.interpolate( x, scale_factor=scale, mode="bilinear", align_corners=True, recompute_scale_factor=True, ) out = self.model(x[0]) res: Dict[str, torch.Tensor] = {} res["boxes"] = out[0] / scale res["labels"] = out[2] res["scores"] = out[1] return inputs, [res] size_divisibility = max(pytorch_model.backbone.size_divisibility, 10) h, w = size_divisibility, size_divisibility * 2 with create_detection_data_loader_on_toy_dataset( cfg, h, w, is_train=False) as data_loader: with make_temp_directory( "test_export_torchvision_format") as tmp_dir: predictor_path = convert_and_export_predictor( cfg, copy.deepcopy(pytorch_model), "torchscript", tmp_dir, data_loader, ) orig_model = torch.jit.load( os.path.join(predictor_path, "model.jit")) wrapped_model = Wrapper(orig_model) # optionally do a forward wrapped_model([torch.rand(3, 600, 600)]) scripted_model = torch.jit.script(wrapped_model) scripted_model.save(os.path.join(tmp_dir, "new_file.pt"))
def test_default_cfg_dump_and_load(self): default_cfg = GeneralizedRCNNRunner().get_default_cfg() cfg = default_cfg.clone() with make_temp_directory("detectron2go_tmp") as tmp_dir: file_name = os.path.join(tmp_dir, "config.yaml") # this is same as the one in fblearner_launch_utils_detectron2go.py with open(file_name, "w") as f: f.write(cfg.dump(default_flow_style=False)) # check if the dumped config file can be merged cfg.merge_from_file(file_name)
def test_load_arch_defs(self): """Test arch def str-to-dict conversion compatible with merging""" default_cfg = GeneralizedRCNNRunner().get_default_cfg() cfg = default_cfg.clone() cfg.merge_from_file(get_resource_path("arch_def_merging.yaml")) with make_temp_directory("detectron2go_tmp") as tmp_dir: # Dump out config with arch def file_name = os.path.join(tmp_dir, "test_archdef_config.yaml") with open(file_name, "w") as f: f.write(cfg.dump()) # Attempt to reload the config another_cfg = default_cfg.clone() another_cfg.merge_from_file(file_name)
def register_toy_coco_dataset( dataset_name, num_images=3, image_size=(5, 10), num_classes=-1, num_keypoints=0 ): width, height = image_size with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir: image_dir = os.path.join(dataset_dir, "images") os.makedirs(image_dir) image_generator = LocalImageGenerator(image_dir, width=width, height=height) with _register_toy_dataset( dataset_name, image_generator, num_images=num_images, num_classes=num_classes, num_keypoints=num_keypoints, ): yield
def test_create_predictor(self): with make_temp_directory("test_model_info") as tmp_dir: # define the predictor model_a_path = os.path.join(tmp_dir, "model_A") predictor_info = PredictorInfo( model=ModelInfo(path=model_a_path, type="torchscript"), preprocess_info=FuncInfo.gen_func_info(TestPreprocess, params={"weight": 2.0}), ) # simulating exporting to predictor _save_test_model(model_a_path) with open(os.path.join(tmp_dir, "predictor_info.json"), "w") as f: json.dump(predictor_info.to_dict(), f) predictor = create_predictor(tmp_dir) # y = (x * 2) + 1 self.assertEqual(torch.tensor(5), predictor(torch.tensor(2)))
def export_optimize_and_save_torchscript( model: nn.Module, inputs: Optional[Tuple[Any]], output_path: str, *, jit_mode: Optional[str] = DEFAULT_JIT_MODE, torchscript_filename: str = "model.jit", mobile_optimization: Optional[MobileOptimizationConfig] = None, _extra_files: Optional[Dict[str, bytes]] = None, ) -> str: """ The primary function for exporting PyTorch model to TorchScript. Args: model (nn.Module): the model to export. When given a ScriptModule, skip the export and only optimize and save model. inputs (tuple or None): input arguments of model, can be called as model(*inputs). Will not be used when scripting the model. output_path (str): directory that the model will be saved. jit_mode (str): trace/script or None if the model is already a ScriptModule. torchscript_filename (str): the filename of non-mobile-optimized model. mobile_optimization (MobileOptimizationConfig): when provided, the mobile optimization will be applied. _extra_files (Dict[str, bytes]): when provided, extra files will be saved. Returns: (str): filename of the final model no matter optmized or not. """ logger.info("Export, optimize and saving TorchScript to {} ...".format( output_path)) PathManager.mkdirs(output_path) if _extra_files is None: _extra_files = {} if isinstance(model, torch.jit.ScriptModule): if jit_mode is not None: logger.info( "The input model is already a ScriptModule, skip the jit step") elif jit_mode == "trace": logger.info("Tracing the model ...") with torch.no_grad(): script_model = torch.jit.trace(model, inputs) elif jit_mode == "script": logger.info("Scripting the model ...") script_model = torch.jit.script(model) else: raise ValueError("Unsupported jit_mode: {}".format(jit_mode)) with make_temp_directory( "export_optimize_and_save_torchscript") as tmp_dir: @contextlib.contextmanager def _synced_local_file(rel_path): remote_file = os.path.join(output_path, rel_path) local_file = os.path.join(tmp_dir, rel_path) yield local_file PathManager.copy_from_local(local_file, remote_file, overwrite=True) with _synced_local_file(torchscript_filename) as model_file: logger.info(f"Saving torchscript model to: {torchscript_filename}") torch.jit.save(script_model, model_file, _extra_files=_extra_files) dump_torchscript_IR(script_model, os.path.join(output_path, "torchscript_IR")) data_filename = "data.pth" with _synced_local_file(data_filename) as data_file: logger.info(f"Saving example data to: {data_filename}") torch.save(inputs, data_file) if mobile_optimization is not None: logger.info("Applying optimize_for_mobile ...") liteopt_model = optimize_for_mobile( script_model, optimization_blocklist=mobile_optimization. optimization_blocklist, preserved_methods=mobile_optimization.preserved_methods, backend=mobile_optimization.backend, ) torchscript_filename = mobile_optimization.torchscript_filename with _synced_local_file(torchscript_filename) as lite_path: logger.info( f"Saving mobile optimized model to: {torchscript_filename}" ) liteopt_model._save_for_lite_interpreter( lite_path, _extra_files=_extra_files) op_names = torch.jit.export_opnames(liteopt_model) logger.info("Operator names from lite interpreter:\n{}".format( "\n".join(op_names))) logger.info("Applying augment_model_with_bundled_inputs ...") # make all tensors zero-like to save storage iters = recursive_iterate(inputs) for x in iters: if isinstance(x, torch.Tensor): iters.send(torch.zeros_like(x).contiguous()) inputs = iters.value augment_model_with_bundled_inputs(liteopt_model, [inputs]) # For non-cpu backends (e.g. Metal, Vulkan) the bundled inputs need to be # converted with `torch.to(<myDevice>)` in order to predict successfully # This is a temporary bypass until PT Edge supports automatic backend # conversion in the bundled inputs interface, or we can auto-add a input tensor # conversion op to Metal and Vulkan models. target_backend = mobile_optimization.backend.lower() if target_backend == "cpu": # Sanity check by running logger.info( "Running sanity check for the mobile optimized model ...") liteopt_model(*liteopt_model.get_all_bundled_inputs()[0]) name, ext = os.path.splitext(torchscript_filename) input_bundled_path = name + "_bundled" + ext with _synced_local_file(input_bundled_path) as lite_path: logger.info( f"Saving input bundled model to: {input_bundled_path}") liteopt_model._save_for_lite_interpreter(lite_path) return torchscript_filename
def trace_and_save_torchscript( model: nn.Module, inputs: Tuple[torch.Tensor], output_path: str, torchscript_filename: str = "model.jit", mobile_optimization: Optional[MobileOptimizationConfig] = None, _extra_files: Optional[Dict[str, bytes]] = None, ): logger.info("Tracing and saving TorchScript to {} ...".format(output_path)) PathManager.mkdirs(output_path) if _extra_files is None: _extra_files = {} with torch.no_grad(): script_model = torch.jit.trace(model, inputs) with make_temp_directory("trace_and_save_torchscript") as tmp_dir: @contextlib.contextmanager def _synced_local_file(rel_path): remote_file = os.path.join(output_path, rel_path) local_file = os.path.join(tmp_dir, rel_path) yield local_file PathManager.copy_from_local(local_file, remote_file, overwrite=True) with _synced_local_file(torchscript_filename) as model_file: torch.jit.save(script_model, model_file, _extra_files=_extra_files) with _synced_local_file("data.pth") as data_file: torch.save(inputs, data_file) if mobile_optimization is not None: logger.info("Applying optimize_for_mobile ...") liteopt_model = optimize_for_mobile( script_model, optimization_blocklist=mobile_optimization. optimization_blocklist, preserved_methods=mobile_optimization.preserved_methods, backend=mobile_optimization.backend, ) torchscript_filename = mobile_optimization.torchscript_filename with _synced_local_file(torchscript_filename) as lite_path: liteopt_model._save_for_lite_interpreter( lite_path, _extra_files=_extra_files) # liteopt_model(*inputs) # sanity check op_names = torch.jit.export_opnames(liteopt_model) logger.info("Operator names from lite interpreter:\n{}".format( "\n".join(op_names))) logger.info("Applying augment_model_with_bundled_inputs ...") # make all tensors zero-like to save storage iters = recursive_iterate(inputs) for x in iters: if isinstance(x, torch.Tensor): iters.send(torch.zeros_like(x).contiguous()) inputs = iters.value augment_model_with_bundled_inputs(liteopt_model, [inputs]) liteopt_model( *liteopt_model.get_all_bundled_inputs()[0]) # sanity check name, ext = os.path.splitext(torchscript_filename) with _synced_local_file(name + "_bundled" + ext) as lite_path: liteopt_model._save_for_lite_interpreter(lite_path) return torchscript_filename
def maskrcnn_export_caffe2_vs_torchvision_opset_format_example(self): with make_temp_directory("export_demo") as tmp_dir: # use a fake dataset for ci dataset_name = create_local_dataset(tmp_dir, 5, 224, 224) config_list = [ "DATASETS.TRAIN", (dataset_name, ), "DATASETS.TEST", (dataset_name, ), ] # START_WIKI_EXAMPLE_TAG runner = GeneralizedRCNNRunner() cfg = runner.get_default_cfg() cfg.merge_from_file("detectron2go://mask_rcnn_fbnetv3a_dsmask_C4.yaml") cfg.merge_from_list(get_quick_test_config_opts()) cfg.merge_from_list(config_list) # equivalent to running: # exporter.par --runner GeneralizedRCNNRunner --config-file config.yaml --predictor-types torchscript tourchscript@c2_ops --output-dir tmp_dir _ = main( cfg, tmp_dir, runner, predictor_types=["torchscript@c2_ops", "torchscript"], ) # the path can be fetched from the return of main, here just use hard-coded values torchvision_ops_model = torch.jit.load( os.path.join(tmp_dir, "torchscript", "model.jit")) caffe2_ops_model = torch.jit.load( os.path.join(tmp_dir, "torchscript@c2_ops", "model.jit")) # Running inference using torchvision-style format image = torch.zeros(1, 64, 96) # chw 3D tensor # The exported model can run on both cpu/gpu device = "cuda:0" if torch.cuda.is_available() else "cpu" torchvision_ops_model = torchvision_ops_model.to(device) torchvision_style_outputs = torchvision_ops_model( image) # suppose N instances are detected # NOTE: the output are flattened tensors of the real output (which is a dict), they're # ordered by the key in dict, which is deterministic for the given model, but it might # be difficult to figure out just from model.jit file. The predictor_info.json from # the same directory contains the `outputs_schema`, which indicate how the final output # is constructed from flattened tensors. ( pred_boxes, # torch.Size([N, 4]) pred_classes, # torch.Size([N]) pred_masks, # torch.Size([N, 1, Hmask, Wmask]) scores, # torch.Size([N]) image_sizes, # torch.Size([2]) ) = torchvision_style_outputs self.assertTrue( all(x.device == torch.device(device) for x in torchvision_style_outputs[:4]), torchvision_style_outputs, ) torch.testing.assert_close(image_sizes, torch.tensor([64, 96])) # Running inference using caffe2-style format data = torch.zeros(1, 1, 64, 96) im_info = torch.tensor([[64, 96, 1.0]]) caffe2_style_outputs = caffe2_ops_model([data, im_info]) # NOTE: the output order is determined in the order of creating the tensor during # forward function, it's also follow the order of original Caffe2 model. roi_bbox_nms = caffe2_style_outputs[0] # torch.Size([N, 4]) roi_score_nms = caffe2_style_outputs[1] # torch.Size([N]) roi_class_nms = caffe2_style_outputs[2] # torch.Size([N]) mask_fcn_probs = caffe2_style_outputs[ 3] # torch.Size([N, Cmask, Hmask, Wmask]) # relations between torchvision-style outputs and caffe2-style outputs torch.testing.assert_close(pred_boxes, roi_bbox_nms, check_device=False) torch.testing.assert_close(pred_classes, roi_class_nms.to(torch.int64), check_device=False) torch.testing.assert_close( pred_masks, mask_fcn_probs[:, roi_class_nms.to(torch.int64), :, :], check_device=False, ) torch.testing.assert_close(scores, roi_score_nms, check_device=False)