def export_tracing(torch_model, inputs): assert TORCH_VERSION >= (1, 8) # RetinaNet is supported but needs a slightly different wrapper. # TODO wrapper should be automatically generated assert isinstance(torch_model, GeneralizedRCNN) image = inputs[0]["image"] class WrapModel(nn.Module): def __init__(self): super().__init__() self.torch_model = torch_model def forward(self, image): inputs = [{"image": image}] outputs = self.torch_model.inference(inputs, do_postprocess=False)[0] outputs = outputs.get_fields() from detectron2.utils.analysis import _flatten_to_tuple return _flatten_to_tuple(outputs) from detectron2.export.torchscript_patch import patch_builtin_len with torch.no_grad(), patch_builtin_len(): assert (args.format == "torchscript" ), "Tracing method only supports torchscript format for now." ts_model = torch.jit.trace(WrapModel(), (image, )) ts_model.save(os.path.join(args.output, "model.ts")) dump_torchscript_IR(ts_model, args.output) # NOTE onnx export fails in pytorch # if args.format == "onnx": # torch.onnx.export(WrapModel(), (image,), os.path.join(args.output, "model.onnx")) # TODO inference in Python now missing postprocessing glue code return None
def test_dump_IR_tracing(self): cfg = get_cfg() cfg.MODEL.RESNETS.DEPTH = 18 cfg.MODEL.RESNETS.RES2_OUT_CHANNELS = 64 class Mod(nn.Module): def forward(self, x): return tuple(self.m(x).values()) model = Mod() model.m = build_backbone(cfg) model.eval() with torch.no_grad(): ts_model = torch.jit.trace(model, (torch.rand(2, 3, 224, 224), )) with tempfile.TemporaryDirectory(prefix="detectron2_test") as d: dump_torchscript_IR(ts_model, d) # check that the files are created for name in [ "model_ts_code", "model_ts_IR", "model_ts_IR_inlined", "model" ]: fname = os.path.join(d, name + ".txt") self.assertTrue(os.stat(fname).st_size > 0, fname)
def export_tracing(torch_model, inputs): assert TORCH_VERSION >= (1, 8) image = inputs[0]["image"] inputs = [{"image": image}] # remove other unused keys if isinstance(torch_model, GeneralizedRCNN): def inference(model, inputs): # use do_postprocess=False so it returns ROI mask inst = model.inference(inputs, do_postprocess=False)[0] return [{"instances": inst}] else: inference = None # assume that we just call the model directly traceable_model = TracingAdapter(torch_model, inputs, inference) if args.format == "torchscript": ts_model = torch.jit.trace(traceable_model, (image, )) with PathManager.open(os.path.join(args.output, "model.ts"), "wb") as f: torch.jit.save(ts_model, f) dump_torchscript_IR(ts_model, args.output) elif args.format == "onnx": # NOTE onnx export currently failing in pytorch with PathManager.open(os.path.join(args.output, "model.onnx"), "wb") as f: torch.onnx.export(traceable_model, (image, ), f) logger.info("Inputs schema: " + str(traceable_model.inputs_schema)) logger.info("Outputs schema: " + str(traceable_model.outputs_schema)) if args.format != "torchscript": return None if not isinstance(torch_model, (GeneralizedRCNN, RetinaNet)): return None def eval_wrapper(inputs): """ The exported model does not contain the final resize step, which is typically unused in deployment but needed for evaluation. We add it manually here. """ input = inputs[0] instances = traceable_model.outputs_schema(ts_model( input["image"]))[0]["instances"] postprocessed = detector_postprocess(instances, input["height"], input["width"]) return [{"instances": postprocessed}] return eval_wrapper
def export_caffe2_tracing(cfg, torch_model, inputs): tracer = Caffe2Tracer(cfg, torch_model, inputs) if args.format == "caffe2": caffe2_model = tracer.export_caffe2() caffe2_model.save_protobuf(args.output) # draw the caffe2 graph caffe2_model.save_graph(os.path.join(args.output, "model.svg"), inputs=inputs) return caffe2_model elif args.format == "onnx": onnx_model = tracer.export_onnx() onnx.save(onnx_model, os.path.join(args.output, "model.onnx")) elif args.format == "torchscript": ts_model = tracer.export_torchscript() ts_model.save(os.path.join(args.output, "model.ts")) dump_torchscript_IR(ts_model, args.output)
def export_scripting(torch_model): assert TORCH_VERSION >= (1, 8) fields = { "proposal_boxes": Boxes, "objectness_logits": Tensor, "pred_boxes": Boxes, "scores": Tensor, "pred_classes": Tensor, "pred_masks": Tensor, "pred_keypoints": torch.Tensor, "pred_keypoint_heatmaps": torch.Tensor, } # maybe can export to onnx format? assert args.format == "torchscript", "Scripting only supports torchscript format." ts_model = export_torchscript_with_instances(torch_model, fields) ts_model.save(os.path.join(args.output, "model.ts")) dump_torchscript_IR(ts_model, args.output) # TODO inference in Python now missing postprocessing glue code return None
data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0]) first_batch = next(iter(data_loader)) # convert and save caffe2 model tracer = Caffe2Tracer(cfg, torch_model, first_batch) if args.format == "caffe2": caffe2_model = tracer.export_caffe2() caffe2_model.save_protobuf(args.output) # draw the caffe2 graph caffe2_model.save_graph(os.path.join(args.output, "model.svg"), inputs=first_batch) elif args.format == "onnx": onnx_model = tracer.export_onnx() onnx.save(onnx_model, os.path.join(args.output, "model.onnx")) elif args.format == "torchscript": ts_model = tracer.export_torchscript() ts_model.save(os.path.join(args.output, "model.ts")) from detectron2.export.torchscript import dump_torchscript_IR dump_torchscript_IR(ts_model, args.output) # run evaluation with the converted model if args.run_eval: assert args.format == "caffe2", "Python inference in other format is not yet supported." dataset = cfg.DATASETS.TEST[0] data_loader = build_detection_test_loader(cfg, dataset) # NOTE: hard-coded evaluator. change to the evaluator for your dataset evaluator = COCOEvaluator(dataset, cfg, True, args.output) metrics = inference_on_dataset(caffe2_model, data_loader, evaluator) print_csv_format(metrics)