def _test_model(self, config_path, inference_func, batch=1): model = model_zoo.get(config_path, trained=True) image = get_sample_coco_image() inputs = tuple(image.clone() for _ in range(batch)) wrapper = TracingAdapter(model, inputs, inference_func) wrapper.eval() with torch.no_grad(): # trace with smaller images, and the trace must still work trace_inputs = tuple( nn.functional.interpolate( image, scale_factor=random.uniform(0.5, 0.7)) for _ in range(batch)) traced_model = torch.jit.trace(wrapper, trace_inputs) outputs = inference_func(model, *inputs) traced_outputs = wrapper.outputs_schema(traced_model(*inputs)) if batch > 1: for output, traced_output in zip(outputs, traced_outputs): assert_instances_allclose(output, traced_output, size_as_tensor=True) else: assert_instances_allclose(outputs, traced_outputs, size_as_tensor=True)
def _test_model(self, config_path, inference_func): model = model_zoo.get(config_path, trained=True) image = get_sample_coco_image() wrapper = TracingAdapter(model, image, inference_func) wrapper.eval() with torch.no_grad(): small_image = nn.functional.interpolate(image, scale_factor=0.5) # trace with a different image, and the trace must still work traced_model = torch.jit.trace(wrapper, (small_image,)) output = inference_func(model, image) traced_output = wrapper.outputs_schema(traced_model(image)) assert_instances_allclose(output, traced_output, size_as_tensor=True)
def test_allow_non_tensor(self): data = (torch.tensor([5, 8]), 3) # contains non-tensor class M(nn.Module): def forward(self, input, number): return input model = M() with self.assertRaisesRegex(ValueError, "must only contain tensors"): adap = TracingAdapter(model, data, allow_non_tensor=False) adap = TracingAdapter(model, data, allow_non_tensor=True) _ = adap(*adap.flattened_inputs) newdata = (data[0].clone(),) with self.assertRaisesRegex(ValueError, "cannot generalize"): _ = adap(*newdata)
def export(cls, model, input_args, save_path, **export_kwargs): adapter = TracingAdapter(model, input_args) trace_and_save_torchscript( adapter, adapter.flattened_inputs, save_path, **export_kwargs ) inputs_schema = dump_dataclass(adapter.inputs_schema) outputs_schema = dump_dataclass(adapter.outputs_schema) return {"inputs_schema": inputs_schema, "outputs_schema": outputs_schema}
def new_f(cls, model, input_args, save_path, export_method, **export_kwargs): force_disable_tracing_adapter = export_kwargs.pop( "force_disable_tracing_adapter", False) is_trace_mode = export_kwargs.get("jit_mode", "trace") == "trace" if force_disable_tracing_adapter or not is_trace_mode: logger.info("Not trace mode, export normally") return old_f(cls, model, input_args, save_path, export_method, **export_kwargs) if _is_data_flattened_tensors(input_args): logger.info( "Dry run the model to check if TracingAdapter is needed ...") outputs = model(*input_args) if _is_data_flattened_tensors(outputs): logger.info( "Both inputs and outputs are flattened tensors, export the model as is." ) load_kwargs = old_f(cls, model, input_args, save_path, export_method, **export_kwargs) assert "tracing_adapted" not in load_kwargs load_kwargs.update({"tracing_adapted": False}) return load_kwargs else: logger.info( "The outputs are not flattened tensors, can't trace normally." ) else: logger.info( "The inputs are not flattened tensors, can't trace normally.") logger.warning( "Wrap the model with TracingAdapter to handle non-flattened inputs/outputs," " please be aware that the exported model will have different input/output data structure." ) adapter = TracingAdapter(model, input_args) load_kwargs = old_f( cls, adapter, adapter.flattened_inputs, save_path, export_method, **export_kwargs, ) inputs_schema = dump_dataclass(adapter.inputs_schema) outputs_schema = dump_dataclass(adapter.outputs_schema) assert "tracing_adapted" not in load_kwargs assert "inputs_schema" not in load_kwargs assert "outputs_schema" not in load_kwargs load_kwargs.update({ "tracing_adapted": True, "inputs_schema": inputs_schema, "outputs_schema": outputs_schema, }) return load_kwargs
def new_f(cls, model, input_args, *args, **kwargs): adapter = TracingAdapter(model, input_args) load_kwargs = old_f(cls, adapter, adapter.flattened_inputs, *args, **kwargs) inputs_schema = dump_dataclass(adapter.inputs_schema) outputs_schema = dump_dataclass(adapter.outputs_schema) assert "inputs_schema" not in load_kwargs assert "outputs_schema" not in load_kwargs load_kwargs.update({ "inputs_schema": inputs_schema, "outputs_schema": outputs_schema }) return load_kwargs
def _test_model(self, config_path, inference_func, batch=1): model = model_zoo.get(config_path, trained=True) image = get_sample_coco_image() inputs = tuple(image.clone() for _ in range(batch)) wrapper = TracingAdapter(model, inputs, inference_func) wrapper.eval() with torch.no_grad(): # trace with smaller images, and the trace must still work trace_inputs = tuple( nn.functional.interpolate(image, scale_factor=random.uniform(0.5, 0.7)) for _ in range(batch) ) traced_model = torch.jit.trace(wrapper, trace_inputs) testing_devices = self._get_device_casting_test_cases(model) # save and load back the model in order to show traceback of TorchScript with tempfile.TemporaryDirectory(prefix="detectron2_test") as d: basename = "model" jitfile = f"{d}/{basename}.jit" torch.jit.save(traced_model, jitfile) traced_model = torch.jit.load(jitfile) if any(device and "cuda" in device for device in testing_devices): self._check_torchscript_no_hardcoded_device(jitfile, d, "cuda") for device in testing_devices: print(f"Testing casting to {device} for inference (traced on {model.device}) ...") with torch.no_grad(): outputs = inference_func(copy.deepcopy(model).to(device), *inputs) traced_outputs = wrapper.outputs_schema(traced_model.to(device)(*inputs)) if batch > 1: for output, traced_output in zip(outputs, traced_outputs): assert_instances_allclose(output, traced_output, size_as_tensor=True) else: assert_instances_allclose(outputs, traced_outputs, size_as_tensor=True)
def d2_meta_arch_prepare_for_export(self, cfg, inputs, predictor_type): if "torchscript" in predictor_type and "@tracing" in predictor_type: def inference_func(model, image): inputs = [{"image": image}] return model.inference(inputs, do_postprocess=False)[0] def data_generator(x): return (x[0]["image"],) image = data_generator(inputs)[0] wrapper = TracingAdapter(self, image, inference_func) wrapper.eval() # HACK: outputs_schema can only be obtained after running tracing, but # PredictorExportConfig requires a pre-defined postprocessing function, this # causes tracing to run twice. logger.info("tracing the model to get outputs_schema ...") with torch.no_grad(), patch_builtin_len(): _ = torch.jit.trace(wrapper, (image,)) outputs_schema_json = json.dumps( wrapper.outputs_schema, default=dataclass_object_dump ) return PredictorExportConfig( model=wrapper, data_generator=data_generator, preprocess_info=FuncInfo.gen_func_info( D2TracingAdapterPreprocessFunc, params={} ), postprocess_info=FuncInfo.gen_func_info( D2TracingAdapterPostFunc, params={"outputs_schema_json": outputs_schema_json}, ), ) if cfg.MODEL.META_ARCHITECTURE in META_ARCH_CAFFE2_EXPORT_TYPE_MAP: C2MetaArch = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[cfg.MODEL.META_ARCHITECTURE] c2_compatible_model = C2MetaArch(cfg, self) preprocess_info = FuncInfo.gen_func_info( D2Caffe2MetaArchPreprocessFunc, params=D2Caffe2MetaArchPreprocessFunc.get_params(cfg, c2_compatible_model), ) postprocess_info = FuncInfo.gen_func_info( D2Caffe2MetaArchPostprocessFunc, params=D2Caffe2MetaArchPostprocessFunc.get_params(cfg, c2_compatible_model), ) preprocess_func = preprocess_info.instantiate() return PredictorExportConfig( model=c2_compatible_model, # Caffe2MetaArch takes a single tuple as input (which is the return of # preprocess_func), data_generator requires all positional args as a tuple. data_generator=lambda x: (preprocess_func(x),), preprocess_info=preprocess_info, postprocess_info=postprocess_info, ) raise NotImplementedError("Can't determine prepare_for_tracing!")