def deploy_nnapi(dataloader: DataLoader, model: nn.Module, fuse: bool, name: str, backend: str = 'qnnpack'): model = deepcopy(model) torch.backends.quantized.engine = backend model.qconfig = torch.quantization.get_default_qconfig(backend) path = f'./{name}_nnapi' model = model.eval() if fuse: model.fuse() path += '_fused' model_prepared = torch.quantization.prepare(model) for sample in dataloader: model_prepared(sample) model_quantized = torch.quantization.convert(model_prepared) input_float = torch.rand(1, 3, 224, 224) quantizer = model_quantized.quant dequantizer = model_quantized.dequant model_quantized.quant = torch.nn.Identity() model_quantized.dequant = torch.nn.Identity() input_tensor = quantizer(input_float) input_tensor = input_tensor.contiguous(memory_format=torch.channels_last) input_tensor.nnapi_nhwc = True with torch.no_grad(): model_quantized_traced = torch.jit.trace(model_quantized, input_tensor) nnapi_model = convert_model_to_nnapi(model_quantized_traced, input_tensor) nnapi_model_float_interface = torch.jit.script( torch.nn.Sequential(quantizer, nnapi_model, dequantizer)) traced_path = path + '_traced.pt' traced_float_path = path + '_float_interface_traced.pt' nnapi_model.save(traced_path) nnapi_model_float_interface.save(traced_float_path)
def check( self, module, arg_or_args, *, trace_args=None, convert_args=None, atol_rtol=None, limit=None, ): with torch.no_grad(): if isinstance(arg_or_args, torch.Tensor): args = [arg_or_args] else: args = arg_or_args module.eval() traced = torch.jit.trace(module, trace_args or args) nnapi_module = convert_model_to_nnapi(traced, convert_args or args) if not self.can_run_nnapi: # Only test that the model was converted successfully. return eager_output = module(*args) nnapi_output = nnapi_module(*args) kwargs = {} if atol_rtol is not None: kwargs["atol"] = atol_rtol[0] kwargs["rtol"] = atol_rtol[1] self.assertEqual(eager_output, nnapi_output, **kwargs) if limit is not None: mismatches = \ eager_output.int_repr().to(torch.int32) - \ nnapi_output.int_repr().to(torch.int32) if mismatches.count_nonzero() > limit: # Too many mismatches. Re-run the check with no tolerance # to get a nice message. self.assertEqual(eager_output, nnapi_output, atol=0, rtol=0)
def call_lowering_to_nnapi(self, traced_module, args): return convert_model_to_nnapi(traced_module, args)