def test_recursive_iter_dict(self): values = {"k1": "v1", "k2": ["v2", 3, "v4"], "k3": 5} val_list = list(iu.recursive_iterate(values)) self.assertEqual(val_list, ["v1", "v2", 3, "v4", 5]) val_list = list(iu.recursive_iterate(values, iter_types=str)) self.assertEqual(val_list, ["v1", "v2", "v4"])
def test_recursive_iter_send(self): values = {"k1": "v1", "k2": ["v2", 3, "v4"], "k3": 5} gt_values = {"k1": "v1_ret", "k2": ["v2_ret", 4, "v4_ret"], "k3": 6} iters = iu.recursive_iterate(values) for x in iters: iters.send(x + "_ret" if isinstance(x, str) else x + 1) result = iters.value self.assertEqual(result, gt_values)
def test_recursive_iter_seq_check_func(self): def _is_int_list(x): return isinstance(x, list) and all(isinstance(y, int) for y in x) values = [{"k1": [1, 2, 3], "k2": ["v2", 3, "v4"], "k3": 5}] # List of ints are not considered as a list val_list = list( iu.recursive_iterate(values, seq_check_func=lambda x: not _is_int_list(x)) ) self.assertEqual(val_list, [[1, 2, 3], "v2", 3, "v4", 5])
def _get_computed_tensor_to_list(op_name, op_cfg_name, result): riter = iu.recursive_iterate(result) for item in riter: item_str = ( f"({list(item.shape)}, " f"{[float('%.5f' % o) for o in item.contiguous().view(-1).tolist()]})" ) riter.send(item_str) ret = f'("{op_name}", "{op_cfg_name}"): {_to_str(riter.value)}, # noqa' return ret
def test_recursive_iter_send_None(self): values = {"k1": "v1", "k2": ["v2", 3, "v4"], "k3": 5} gt_values = { "k1": "v1_ret", "k2": ["v2_ret", None, "v4_ret"], "k3": None, } iters = iu.recursive_iterate(values, wait_on_send=True) for x in iters: iters.send(x + "_ret" if isinstance(x, str) else None) result = iters.value self.assertEqual(result, gt_values)
def build_upsample(name=None, scales=None, **kwargs): if name is None or scales is None: return None if all(x == 1 for x in iu.recursive_iterate(scales)): return None if name == "default": ret = Upsample(scale_factor=scales, **kwargs) else: ret = UPSAMPLE_REGISTRY.get(name)(scales, **kwargs) return ret
def post_training_quantize(cfg, model, data_loader): """Calibrate a model, convert it to a quantized pytorch model""" model = copy.deepcopy(model) model.eval() # TODO: check why some parameters will have gradient for param in model.parameters(): param.requires_grad = False if hasattr(model, "prepare_for_quant"): model = model.prepare_for_quant(cfg) else: logger.info("Using default implementation for prepare_for_quant") model = default_prepare_for_quant(cfg, model) if cfg.QUANTIZATION.EAGER_MODE: torch.quantization.prepare(model, inplace=True) logger.info("Prepared the PTQ model for calibration:\n{}".format(model)) # Option for forcing running calibration on GPU, works only when the model supports # casting both model and inputs. calibration_force_on_gpu = ( cfg.QUANTIZATION.PTQ.CALIBRATION_FORCE_ON_GPU and torch.cuda.is_available() ) if calibration_force_on_gpu: # NOTE: model.to(device) may not handle cases such as normalizer, FPN, only # do move to GPU if specified. _cast_detection_model(model, "cuda") calibration_iters = cfg.QUANTIZATION.PTQ.CALIBRATION_NUM_IMAGES for idx, inputs in enumerate(data_loader): logger.info("Running calibration iter: {}/{}".format(idx, calibration_iters)) if calibration_force_on_gpu: iters = recursive_iterate(inputs) for x in iters: if isinstance(x, torch.Tensor): iters.send(x.to("cuda")) inputs = iters.value with torch.no_grad(): model(inputs) if idx + 1 == calibration_iters: break else: logger.warning("Can't run enough calibration iterations") # cast model back to the original device if calibration_force_on_gpu: _cast_detection_model(model, cfg.MODEL.DEVICE) return model
def test_recursive_iter_map_check_func(self): def _is_int_dict_key(x): return isinstance(x, dict) and all(isinstance(yk, int) for yk in x.keys()) values = [{"k1": {1: 1, 2: 2, 3: 3}, "k2": ["v2", 3, "v4"], "k3": 5}] # dict where keys are ints are not considered as a map val_list = list( iu.recursive_iterate( values, iter_types=dict, map_check_func=lambda x: not _is_int_dict_key(x), ) ) self.assertEqual(val_list, [{1: 1, 2: 2, 3: 3}])
def test_recursive_iter_simple(self): self.assertEqual(next(iter(iu.recursive_iterate(1))), 1) self.assertEqual(next(iter(iu.recursive_iterate("str"))), "str") self.assertEqual(list(iu.recursive_iterate(1)), [1]) self.assertEqual(list(iu.recursive_iterate("str")), ["str"]) # special cases for empty list and dict, the function will return an # empty iterable self.assertEqual(list(iu.recursive_iterate([])), []) self.assertEqual(list(iu.recursive_iterate({})), [])
def _get_expected_output_to_tensor(outputs): # a valid tensor is represented by a tuple with two lists (shape and value) def _is_shape_value_tuple(obj): if isinstance(obj, tuple) and len(obj) == 2: if isinstance(obj[0], list) and isinstance(obj[1], list): return True return False # get all tensors riter = iu.recursive_iterate( outputs, seq_check_func=lambda x: not _is_shape_value_tuple(x) ) for item in riter: item_tensor = torch.FloatTensor(item[1]).reshape(item[0]) riter.send(item_tensor) # get all tensor shapes shape_iter = iu.recursive_iterate( outputs, seq_check_func=lambda x: not _is_shape_value_tuple(x) ) for item in shape_iter: shape_iter.send(item[0]) return riter.value, shape_iter.value
def convert_torch_script( model, inputs, fuse_bn=True, verify_output=True, use_get_traceable=False ): assert isinstance(inputs, (tuple, list)), f"Invalid input types {inputs}" if verify_output: print("Run pytorch model") with torch.no_grad(): output_before = model(*inputs) if fuse_bn: print("Fusing bn...") fused_model = fuse_utils.fuse_model(model) if fuse_utils.check_bn_exist(fused_model): print(f"WARNING: BN existed after fusing, {fused_model}") else: fused_model = copy.deepcopy(model) for x in fused_model.parameters(): x.requires_grad = False if use_get_traceable: print("Get traceable model...") fused_model = ju.get_traceable_model(fused_model) print("Start tracing...") with torch.no_grad(): traced_model = torch.jit.trace(fused_model, inputs, strict=False) # print(f"Traced model {traced_model}") # print(f"Traced model {traced_model.code}") # print("Optimizing traced model...") # traced_model = optimize_for_mobile(traced_model) print("Generating traced model lints...") print(generate_mobile_module_lints(traced_model)) print("Run traced model") with torch.no_grad(): outputs = traced_model(*inputs) if verify_output: paired_outputs = iu.create_pair(output_before, outputs) for x in iu.recursive_iterate(paired_outputs, iter_types=torch.Tensor): np.testing.assert_allclose( x.lhs.detach(), x.rhs.detach(), rtol=0, atol=1e-4 ) return traced_model, outputs
def test_paired(self): self.assertIsInstance(iu.PairedDict({}, {}), iu.cabc.Mapping) self.assertIsInstance(iu.PairedSeq([], []), iu.cabc.Sequence) lhs = {"alist": [1, 2, 3], "bdict": {"c": "d", "e": ["f", "g", "h"]}} rhs = {"alist": ["1", "2", "3"], "bdict": {"c": 4, "e": [5, 6, 7]}} paired = iu.create_pair(lhs, rhs) self.assertIsInstance(paired, iu.PairedDict) paired_list = paired["alist"] self.assertIsInstance(paired_list, iu.PairedSeq) paired_obj = paired_list[0] self.assertIsInstance(paired_obj, iu.Pair) iter = iu.recursive_iterate(paired) for x in iter: iter.send(x.to_tuple()) merged_gt = { "alist": [(1, "1"), (2, "2"), (3, "3")], "bdict": {"c": ("d", 4), "e": [("f", 5), ("g", 6), ("h", 7)]}, } self.assertEqual(iter.value, merged_gt)
def _compare_output_shape(self, outputs, gt_shapes): for item in iu.recursive_iterate(iu.create_pair(outputs, gt_shapes)): self.assertEqual(item.lhs.shape, torch.Size(item.rhs))
def get_neg_stride(stride): iters = iu.recursive_iterate(stride) for ss in iters: assert ss is not None iters.send(-ss) return iters.value
def is_neg_stride(stride): return all(x < 0 for x in iu.recursive_iterate(stride))
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 _compare_outputs(outputs, gt_outputs, error_msg): riter = iu.recursive_iterate(iu.create_pair(outputs, gt_outputs)) for item in riter: np.testing.assert_allclose( item.lhs, item.rhs, rtol=0, atol=1e-4, err_msg=error_msg )
def test_recursive_iter_simple_no_wait_on_send(self): self.assertEqual(list(iu.recursive_iterate(None, wait_on_send=False)), [None]) self.assertEqual( list(iu.recursive_iterate([1, 2, None], wait_on_send=False)), [1, 2, None], )
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 forward(self, x): assert len(list(iu.recursive_iterate(x))) == len(self.stubs) data_iter = iu.recursive_iterate(x, wait_on_send=True) for data, stub in zip(data_iter, self.stubs): data_iter.send(stub(data)) return data_iter.value