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, data_loader): inputs = next(iter(data_loader)) images = torch_model.preprocess_image(inputs) images_tensor = images.tensor # If you want Run with FP16, 'images.tensor.half()' instead 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, images_tensor, inference) ts_model = torch.jit.trace(traceable_model, (images_tensor,)) 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) logger.info("Inputs schema: " + str(traceable_model.inputs_schema)) logger.info("Outputs schema: " + str(traceable_model.outputs_schema)) return ts_model
def test_dump_IR_function(self): @torch.jit.script def gunc(x, y): return x + y def func(x, y): return x + y + gunc(x, y) ts_model = torch.jit.trace(func, (torch.rand(3), torch.rand(3))) with tempfile.TemporaryDirectory(prefix="detectron2_test") as d: dump_torchscript_IR(ts_model, d) for name in ["model_ts_code", "model_ts_IR", "model_ts_IR_inlined"]: 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, opset_version=11) 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_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, } assert args.format == "torchscript", "Scripting only supports torchscript format." class ScriptableAdapterBase(nn.Module): # Use this adapter to workaround https://github.com/pytorch/pytorch/issues/46944 # by not retuning instances but dicts. Otherwise the exported model is not deployable def __init__(self): super().__init__() self.model = torch_model self.eval() if isinstance(torch_model, GeneralizedRCNN): class ScriptableAdapter(ScriptableAdapterBase): def forward( self, inputs: Tuple[Dict[str, torch.Tensor]]) -> List[Dict[str, Tensor]]: instances = self.model.inference(inputs, do_postprocess=False) return [i.get_fields() for i in instances] else: class ScriptableAdapter(ScriptableAdapterBase): def forward( self, inputs: Tuple[Dict[str, torch.Tensor]]) -> List[Dict[str, Tensor]]: instances = self.model(inputs) return [i.get_fields() for i in instances] ts_model = scripting_with_instances(ScriptableAdapter(), fields) 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) # TODO inference in Python now missing postprocessing glue code return None
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() 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)
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, } assert args.format == "torchscript", "Scripting only supports torchscript format." ts_model = scripting_with_instances(torch_model, fields) 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) # TODO inference in Python now missing postprocessing glue code return None
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