| #!/usr/bin/env python3 | |
| # Owner(s): ["oncall: mobile"] | |
| import os | |
| import ctypes | |
| import torch | |
| import unittest | |
| from typing import Tuple | |
| from torch.backends._nnapi.prepare import convert_model_to_nnapi | |
| from torch.testing._internal.common_quantized import supported_qengines | |
| from torch.testing._internal.common_utils import TestCase, run_tests | |
| def qpt(t, scale, zero_point, dtype=torch.quint8): | |
| t = torch.tensor(t) | |
| return torch.quantize_per_tensor(t, scale, zero_point, dtype) | |
| def nhwc(t): | |
| t = t.clone().contiguous(memory_format=torch.channels_last) | |
| t.nnapi_nhwc = True | |
| return t | |
| @unittest.skipUnless('qnnpack' in supported_qengines, | |
| "This Pytorch Build has not been built with or does not support QNNPACK") | |
| class TestNNAPI(TestCase): | |
| def setUp(self): | |
| # Avoid saturation in fbgemm | |
| torch.backends.quantized.engine = 'qnnpack' | |
| libneuralnetworks_path = os.environ.get("LIBNEURALNETWORKS_PATH") | |
| if libneuralnetworks_path: | |
| ctypes.cdll.LoadLibrary(libneuralnetworks_path) | |
| print("Will attempt to run NNAPI models.") | |
| self.can_run_nnapi = True | |
| else: | |
| self.can_run_nnapi = False | |
| # Created for easy override by subclasses (eg TestNnapiBackend) | |
| def call_lowering_to_nnapi(self, traced_module, args): | |
| return convert_model_to_nnapi(traced_module, args) | |
| # Created for subclasses to set can_run_nnapi (eg TestNnapiBackend) | |
| def set_can_run_nnapi(self, can_run): | |
| self.can_run_nnapi = can_run | |
| def check( | |
| self, | |
| module, | |
| arg_or_args, | |
| *, | |
| trace_args=None, | |
| convert_args=None, | |
| atol_rtol=None, | |
| limit=None, | |
| expected_memory_format=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 = self.call_lowering_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) | |
| if expected_memory_format: | |
| self.assertTrue(nnapi_output.is_contiguous(memory_format=expected_memory_format)) | |
| def float_and_quant_and_nhwc(self, inp_float, scale, zero_point): | |
| torch.manual_seed(29) | |
| inp_quant = qpt(inp_float, 0.03, 128) | |
| return [ | |
| ("float", inp_float), | |
| ("float-nhwc", nhwc(inp_float)), | |
| ("quant", inp_quant), | |
| ("quant-nhwc", nhwc(inp_quant)), | |
| ] | |
| def test_prelu(self): | |
| arg = torch.tensor([[1.0, -1.0, 2.0, -2.0]]).unsqueeze(-1).unsqueeze(-1) | |
| single_a = torch.nn.PReLU() | |
| self.check(single_a, arg) | |
| multi_a = torch.nn.PReLU(4) | |
| with torch.no_grad(): | |
| multi_a.weight.copy_(torch.tensor([.1, .2, .3, .4])) | |
| self.check(multi_a, nhwc(arg)) | |
| # Test flexible size | |
| self.check( | |
| multi_a, | |
| arg, | |
| trace_args=[torch.zeros(1, 4, 3, 3)], | |
| convert_args=[nhwc(torch.zeros(1, 4, 0, 0))], | |
| ) | |
| def test_quantize(self): | |
| self.check( | |
| torch.ao.nn.quantized.Quantize(0.25, 2, torch.quint8), | |
| nhwc(torch.tensor([[[[1.0]], [[2.0]]]]))) | |
| def test_dequantize(self): | |
| self.check( | |
| torch.ao.nn.quantized.DeQuantize(), | |
| nhwc(qpt([[[[1.0]], [[2.0]]]], 0.25, 2))) | |
| def test_unsqueeze(self): | |
| class UnsqueezeModule(torch.nn.Module): | |
| def __init__(self, dim): | |
| super().__init__() | |
| self.dim = dim | |
| def forward(self, arg): | |
| return arg.unsqueeze(self.dim) | |
| self.check(UnsqueezeModule(-2), torch.randn(4, 2, 2)) | |
| self.check(UnsqueezeModule(-1), torch.randn(4, 2, 2)) | |
| self.check(UnsqueezeModule(0), torch.randn(4, 2, 2)) | |
| self.check(UnsqueezeModule(1), torch.randn(4, 2, 2)) | |
| self.check(UnsqueezeModule(2), torch.randn(4, 2, 2)) | |
| def test_reshape(self): | |
| class ReshapeModule(torch.nn.Module): | |
| def __init__(self, shape): | |
| super().__init__() | |
| self.shape = shape | |
| def forward(self, arg): | |
| return arg.reshape(self.shape) | |
| self.check( | |
| ReshapeModule((2, 4)), | |
| torch.randn(4, 2, 1, 1)) | |
| self.check( | |
| ReshapeModule((8, -1)), | |
| nhwc(torch.randn(4, 2, 1, 1))) | |
| with self.assertRaisesRegex(Exception, "target size"): | |
| self.check( | |
| ReshapeModule((2, 4)), | |
| nhwc(torch.randn(4, 2, 1, 1))) | |
| def test_flatten(self): | |
| for mod in [ | |
| torch.nn.Flatten(), | |
| torch.nn.Flatten(start_dim=2, end_dim=3), | |
| torch.nn.Flatten(start_dim=2, end_dim=4), | |
| torch.nn.Flatten(start_dim=0, end_dim=-2), | |
| torch.nn.Flatten(start_dim=0, end_dim=4) | |
| ]: | |
| self.check(mod, torch.randn(4, 2, 1, 3, 7)) | |
| # flex inputs | |
| self.check( | |
| torch.nn.Flatten(), | |
| torch.randn(4, 2, 1, 3, 7), | |
| convert_args=[torch.zeros(0, 2, 1, 3, 7)] | |
| ) | |
| # channels last | |
| self.check( | |
| torch.nn.Flatten(), | |
| nhwc(torch.randn(2, 1, 4, 7)) | |
| ) | |
| self.check( | |
| torch.nn.Flatten(), | |
| nhwc(torch.randn(2, 3, 1, 1)) | |
| ) | |
| # Exceptions | |
| with self.assertRaisesRegex(Exception, "not supported on NHWC"): | |
| self.check( | |
| torch.nn.Flatten(), | |
| nhwc(torch.randn(1, 3, 4, 4)) | |
| ) | |
| with self.assertRaisesRegex(Exception, "Flattening flexible dims is not supported yet"): | |
| self.check(torch.nn.Flatten(), torch.randn(4, 2, 0, 0, 7)) | |
| with self.assertRaisesRegex(Exception, "Only 1 dim"): | |
| self.check( | |
| torch.nn.Flatten(start_dim=1, end_dim=-2), | |
| torch.randn(0, 2, 1, 3, 0)) | |
| def test_slice(self): | |
| class SliceModule(torch.nn.Module): | |
| def __init__(self, start, stop, step): | |
| super().__init__() | |
| self.start = start | |
| self.stop = stop | |
| self.step = step | |
| def forward(self, t): | |
| return t[1:, self.start:self.stop:self.step, :] | |
| class SliceModule2(torch.nn.Module): | |
| def forward(self, t): | |
| return t[3:] | |
| self.check( | |
| SliceModule(1, 5, 2), | |
| torch.randn(4, 6, 2) | |
| ) | |
| self.check( | |
| SliceModule2(), | |
| torch.randn(5) | |
| ) | |
| # flex inputs | |
| self.check( | |
| SliceModule(1, 5, 2), | |
| torch.randn(4, 6, 2), | |
| convert_args=[torch.zeros(4, 6, 0)] | |
| ) | |
| with self.assertRaisesRegex(Exception, "slice with flexible shape"): | |
| self.check( | |
| SliceModule(1, 5, 2), | |
| torch.randn(4, 6, 2), | |
| convert_args=[torch.zeros(0, 0, 0)] | |
| ) | |
| def test_cat(self): | |
| class CatModule(torch.nn.Module): | |
| def __init__(self, dim): | |
| super().__init__() | |
| self.dim = dim | |
| def forward(self, t1, t2): | |
| return torch.cat([t1, t2], self.dim) | |
| self.check( | |
| CatModule(0), | |
| [ | |
| torch.randn(1, 2, 3, 3), | |
| torch.randn(2, 2, 3, 3), | |
| ]) | |
| self.check( | |
| CatModule(1), | |
| [ | |
| torch.randn(1, 2, 3, 3), | |
| torch.randn(1, 4, 3, 3), | |
| ]) | |
| self.check( | |
| CatModule(1), | |
| [ | |
| nhwc(torch.randn(1, 2, 3, 3)), | |
| nhwc(torch.randn(1, 4, 3, 3)), | |
| ]) | |
| self.check( | |
| CatModule(1), | |
| [ | |
| torch.randn(1, 2, 3, 3), | |
| torch.randn(1, 4, 3, 3), | |
| ], | |
| convert_args=[ | |
| torch.zeros(0, 0, 0, 0), | |
| torch.zeros(0, 0, 0, 0) | |
| ]) | |
| def test_pointwise_unary(self): | |
| for op in ["relu", "sigmoid"]: | |
| with self.subTest(op): | |
| class UnaryModule(torch.nn.Module): | |
| def forward(self, arg): | |
| if op == "relu": | |
| return torch.nn.functional.relu(arg) | |
| if op == "sigmoid": | |
| return torch.sigmoid(arg) | |
| raise Exception("Bad op") | |
| self.check(UnaryModule(), torch.tensor([-1.0, 1.0])) | |
| self.check( | |
| UnaryModule(), | |
| qpt(torch.tensor([-1.0, 1.0]), 1. / 256, 0), | |
| ) | |
| def test_pointwise_binary(self): | |
| for op in ["add", "sub", "mul", "div"]: | |
| with self.subTest(op): | |
| class BinaryModule(torch.nn.Module): | |
| def forward(self, lhs, rhs): | |
| if op == "add": | |
| return lhs + rhs | |
| if op == "sub": | |
| return lhs - rhs | |
| if op == "mul": | |
| return lhs * rhs | |
| if op == "div": | |
| return lhs / rhs | |
| raise Exception("Bad op") | |
| self.check( | |
| BinaryModule(), | |
| [ | |
| torch.tensor([1.0, 2.0]), | |
| torch.tensor([3.0, 4.0]), | |
| ]) | |
| self.check( | |
| BinaryModule(), | |
| [ | |
| torch.tensor([[1.0, 2.0]]), | |
| torch.tensor([[3.0, 4.0], [5.0, 6.0]]), | |
| ]) | |
| with self.assertRaisesRegex(Exception, "Non-equal-rank broadcast"): | |
| self.check( | |
| BinaryModule(), | |
| [ | |
| torch.tensor([1.0, 2.0]), | |
| torch.tensor([[3.0, 4.0], [5.0, 6.0]]), | |
| ]) | |
| def test_pointwise_binary_const(self): | |
| const = torch.randn(1, 4, 6, 6) | |
| class ArgPlusConst(torch.nn.Module): | |
| def forward(self, arg): | |
| return arg + const | |
| class ConstPlusArg(torch.nn.Module): | |
| def forward(self, arg): | |
| return const + arg | |
| arg_contig = torch.randn(2, 4, 6, 6) | |
| arg_nhwc = nhwc(torch.randn(2, 4, 6, 6)) | |
| for mod_class in [ArgPlusConst, ConstPlusArg]: | |
| for use_nhwc in [False, True]: | |
| with self.subTest(mod_class=mod_class.__name__, use_nhwc=use_nhwc): | |
| arg = arg_nhwc if use_nhwc else arg_contig | |
| memory_format = torch.channels_last if use_nhwc else torch.contiguous_format | |
| self.check(mod_class(), arg, | |
| expected_memory_format=memory_format) | |
| def test_hardtanh(self): | |
| inp = torch.tensor([-2.0, -0.5, 0.5, 2.0, 7.0]) | |
| self.check(torch.nn.Hardtanh(), inp) | |
| self.check(torch.nn.Hardtanh(0.0, 6.0), inp) | |
| with self.assertRaisesRegex(Exception, "hardtanh with args"): | |
| self.check(torch.nn.Hardtanh(0.0, 5.0), inp) | |
| def test_softmax(self): | |
| inp = torch.tensor([[-2.0, -0.5], [0.5, 2.0]]) | |
| self.check(torch.nn.Softmax(), inp) | |
| self.check(torch.nn.Softmax(dim=0), inp) | |
| # Test flexible size | |
| self.check( | |
| torch.nn.Softmax(), | |
| inp, | |
| convert_args=[torch.zeros(0, 0)], | |
| ) | |
| def test_to(self): | |
| class ToCPU(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.prelu = torch.nn.PReLU() | |
| def forward(self, x): | |
| y = x.to("cpu") | |
| # add prelu since input operand can't be output | |
| return self.prelu(y) | |
| arg = torch.randn(1, 2, 3, 3) | |
| self.check(ToCPU(), arg) | |
| # Test flexible size | |
| self.check( | |
| ToCPU(), | |
| arg, | |
| convert_args=[torch.zeros(1, 2, 0, 0)], | |
| ) | |
| def test_detach(self): | |
| class DetachModule(torch.nn.Module): | |
| def forward(self, x): | |
| y = x.detach() | |
| return torch.nn.functional.relu(y) | |
| self.check(DetachModule(), torch.randn(1, 2, 3, 3)) | |
| self.check( | |
| DetachModule(), torch.randn(1, 2, 3, 3), | |
| convert_args=[torch.zeros(1, 2, 0, 0)]) | |
| def test_log_softmax(self): | |
| inp = torch.randn(3, 10) | |
| self.check(torch.nn.LogSoftmax(), inp) | |
| self.check(torch.nn.LogSoftmax(0), inp) | |
| def test_mean(self): | |
| class MeanModule(torch.nn.Module): | |
| def __init__(self, dim, keep=False): | |
| super().__init__() | |
| self.dim = dim | |
| self.keep = keep | |
| def forward(self, t): | |
| return torch.mean(t, dim=self.dim, keepdim=self.keep) | |
| self.check(MeanModule(0), torch.randn(2, 3)) | |
| self.check(MeanModule(1), torch.randn(2, 3)) | |
| self.check(MeanModule([2, 3]), torch.randn(2, 3, 6, 6)) | |
| self.check(MeanModule([2, 3]), nhwc(torch.randn(2, 3, 6, 6))) | |
| self.check(MeanModule([-1, -2]), nhwc(torch.randn(2, 3, 6, 6))) | |
| self.check(MeanModule([-1, -2], keep=True), nhwc(torch.randn(2, 3, 6, 6))) | |
| def test_max_pool2d(self): | |
| for (name, inp) in self.float_and_quant_and_nhwc(torch.randn(2, 3, 12, 16), 0.3, 128): | |
| with self.subTest(name): | |
| self.check(torch.nn.MaxPool2d(2), inp) | |
| self.check(torch.nn.MaxPool2d((3, 4)), inp) | |
| self.check(torch.nn.MaxPool2d((3, 4), (1, 2)), inp) | |
| def test_avg_pool2d(self): | |
| for (name, inp) in self.float_and_quant_and_nhwc(torch.randn(2, 3, 12, 16), 0.3, 128): | |
| with self.subTest(name): | |
| atol_rtol = None | |
| limit = None | |
| convert_dims = (2, 3, 0, 0) | |
| convert_arg = torch.zeros(*convert_dims) | |
| for model in ( | |
| torch.nn.AvgPool2d(2), | |
| torch.nn.AvgPool2d((3, 4)), | |
| torch.nn.AvgPool2d((3, 4), (1, 2))): | |
| if "quant" in name: | |
| atol_rtol = (1, 0) | |
| limit = model(inp).numel() | |
| convert_arg = qpt(torch.zeros(*convert_dims), 1.0 / 16, 128) | |
| if "nhwc" in name: | |
| convert_arg = nhwc(convert_arg) | |
| self.check(model, inp, atol_rtol=atol_rtol, limit=limit) | |
| self.check( | |
| model, | |
| inp, | |
| convert_args=[convert_arg], | |
| atol_rtol=atol_rtol, | |
| limit=limit | |
| ) | |
| def test_adaptive_avg_pool2d(self): | |
| for (name, inp) in self.float_and_quant_and_nhwc(torch.randn(2, 3, 12, 16), 0.3, 128): | |
| with self.subTest(name): | |
| self.check(torch.nn.AdaptiveAvgPool2d((1, 1)), inp) | |
| with self.assertRaisesRegex(Exception, "with output size"): | |
| self.check(torch.nn.AdaptiveAvgPool2d((2, 2)), inp) | |
| def test_upsample_nearest2d(self): | |
| convert_args = dict(self.float_and_quant_and_nhwc(torch.randn(2, 3, 0, 0), 0.3, 128)) | |
| for (name, inp) in self.float_and_quant_and_nhwc(torch.randn(2, 3, 12, 16), 0.3, 128): | |
| with self.subTest(name): | |
| self.check(torch.nn.UpsamplingNearest2d(size=(16, 20)), inp) | |
| self.check(torch.nn.UpsamplingNearest2d(size=(24, 32)), inp) | |
| self.check(torch.nn.UpsamplingNearest2d(size=(36, 48)), inp) | |
| self.check(torch.nn.UpsamplingNearest2d(scale_factor=(1.5, 1.5)), inp) | |
| self.check(torch.nn.UpsamplingNearest2d(scale_factor=(2.0, 2.0)), inp) | |
| self.check(torch.nn.UpsamplingNearest2d(scale_factor=(3.0, 3.0)), inp) | |
| self.check( | |
| torch.nn.UpsamplingNearest2d(size=(24, 32)), inp, | |
| convert_args=[convert_args[name]] | |
| ) | |
| self.check( | |
| torch.nn.UpsamplingNearest2d(scale_factor=(2.0, 2.0)), inp, | |
| convert_args=[convert_args[name]] | |
| ) | |
| def test_linear(self): | |
| torch.manual_seed(29) | |
| self.check(torch.nn.Linear(16, 32), torch.randn(2, 16)) | |
| self.check( | |
| torch.nn.Linear(16, 32), torch.randn(2, 16), | |
| convert_args=[torch.zeros(0, 16)]) | |
| def test_conv2d(self): | |
| cases = [ | |
| # in_ch, out_ch, kernel, stride, padding, groups, bias, input_dim, name | |
| ( 4, 8, (3, 3), 1, 0, 1, 1, (2, 4, 16, 16), "3x3"), # noqa: E201,E241 | |
| ( 4, 8, (3, 3), 1, 0, 1, 0, (2, 4, 16, 16), "3x3nobias"), # noqa: E201,E241 | |
| ( 4, 16, (3, 3), 1, 1, 1, 1, (2, 4, 16, 16), "3x3p1"), # noqa: E201,E241 | |
| ( 8, 8, (3, 3), 2, 0, 1, 1, (2, 8, 16, 16), "3x3s2"), # noqa: E201,E241 | |
| ( 4, 8, (5, 5), 1, 0, 1, 1, (2, 4, 16, 16), "5x5"), # noqa: E201,E241 | |
| ( 4, 4, (3, 3), 1, 0, 4, 1, (2, 4, 16, 16), "3x3dw"), # noqa: E201,E241 | |
| ( 8, 4, (1, 1), 1, 0, 1, 1, (2, 8, 16, 16), "1x1"), # noqa: E201,E241 | |
| ] | |
| for kind in ["float", "float-nhwc", "quant", "quant-nhwc"]: | |
| for case in cases: | |
| in_ch, out_ch, kernel, stride, padding, groups, bias, input_dim, name = case | |
| with self.subTest(f"{kind}-{name}"): | |
| inp = torch.randn(input_dim) | |
| model = torch.nn.Conv2d(in_ch, out_ch, kernel, stride, padding, groups=groups, bias=bool(bias)) | |
| output_size = model(inp).numel() | |
| atol_rtol = None | |
| limit = None | |
| convert_dims = (0, in_ch, 0, 0) | |
| convert_arg = torch.zeros(*convert_dims) | |
| if "quant" in kind: | |
| model = torch.nn.Sequential(model) | |
| model.eval() | |
| model.qconfig = torch.ao.quantization.get_default_qconfig('qnnpack') | |
| model = torch.ao.quantization.prepare(model) | |
| model(inp) | |
| model = torch.ao.quantization.convert(model) | |
| inp = qpt(inp, 1.0 / 16, 128) | |
| # I've seen numerical differences between QNNPACK and NNAPI, | |
| # but never more than 1 quantum, and never more than ~1% of | |
| # the output in this test. | |
| atol_rtol = (1, 0) | |
| limit = output_size * 0.03 | |
| convert_arg = qpt(torch.zeros(*convert_dims), 1.0 / 16, 128) | |
| if "nhwc" in kind: | |
| inp = nhwc(inp) | |
| convert_arg = nhwc(convert_arg) | |
| self.check(model, inp, atol_rtol=atol_rtol, limit=limit) | |
| self.check( | |
| model, | |
| inp, | |
| convert_args=[convert_arg], | |
| atol_rtol=atol_rtol, | |
| limit=limit | |
| ) | |
| def test_conv2d_transpose(self): | |
| torch.manual_seed(29) | |
| in_ch, out_ch, kernel = (5, 7, (2, 2)) | |
| input_dim = (4, 5, 3, 3) | |
| convert_dims = input_dim[:2] + (0, 0) | |
| for kind in ["float", "float-nhwc", "quant", "quant-nhwc"]: | |
| with self.subTest(kind): | |
| inp = torch.randn(input_dim) | |
| model = torch.nn.ConvTranspose2d(in_ch, out_ch, kernel) | |
| output_size = model(inp).numel() | |
| atol_rtol = (0.0002, 0) | |
| limit = None | |
| convert_arg = torch.zeros(*convert_dims) | |
| if "quant" in kind: | |
| model = torch.ao.nn.quantized.ConvTranspose2d(in_ch, out_ch, kernel) | |
| model.qconfig = torch.ao.quantization.get_default_qconfig('qnnpack') | |
| inp = qpt(inp, 1.0 / 16, 128) | |
| # I've seen numerical differences between QNNPACK and NNAPI, | |
| # but never more than 1 quantum, and never more than ~10% of | |
| # the output in this test. | |
| atol_rtol = (1, 0) | |
| limit = output_size * 0.1 | |
| convert_arg = qpt(convert_arg, 1.0 / 16, 128) | |
| if "nhwc" in kind: | |
| inp = nhwc(inp) | |
| convert_arg = nhwc(convert_arg) | |
| self.check(model, inp, atol_rtol=atol_rtol, limit=limit) | |
| self.check( | |
| model, | |
| inp, | |
| convert_args=[convert_arg], | |
| atol_rtol=atol_rtol, | |
| limit=limit | |
| ) | |
| def test_qadd(self): | |
| func = torch.ao.nn.quantized.QFunctional() | |
| func.scale = 0.5 | |
| func.zero_point = 120 | |
| class AddMod(torch.nn.Module): | |
| def forward(self, lhs, rhs): | |
| return func.add(lhs, rhs) | |
| class AddReluMod(torch.nn.Module): | |
| def forward(self, lhs, rhs): | |
| return func.add_relu(lhs, rhs) | |
| class MulMod(torch.nn.Module): | |
| def forward(self, lhs, rhs): | |
| return func.mul(lhs, rhs) | |
| for (name, mod) in [("add", AddMod), ("add_relu", AddReluMod), ("mul", MulMod)]: | |
| with self.subTest(name): | |
| self.check( | |
| mod(), | |
| [ | |
| qpt([1.0, 2.0], 0.25, 128), | |
| qpt([3.0, 4.0], 0.25, 128), | |
| ]) | |
| self.check( | |
| mod(), | |
| [ | |
| qpt([[1.0, 2.0]], 0.25, 128), | |
| qpt([[3.0, 4.0]], 0.25, 128), | |
| ], | |
| convert_args=[ | |
| qpt([[1.0, 2.0]], 0.25, 128), | |
| qpt(torch.zeros((1, 2)), 0.25, 128), | |
| ] | |
| ) | |
| self.check( | |
| mod(), | |
| [ | |
| qpt([[1.0, 2.0]], 0.25, 128), | |
| qpt([[3.0, 4.0]], 0.25, 128), | |
| ], | |
| convert_args=[ | |
| qpt(torch.zeros((1, 2)), 0.25, 128), | |
| qpt([[3.0, 4.0]], 0.25, 128), | |
| ] | |
| ) | |
| self.check( | |
| mod(), | |
| [ | |
| qpt([[1.0, 2.0]], 0.25, 128), | |
| qpt([[3.0, 4.0]], 0.25, 128), | |
| ], | |
| convert_args=[ | |
| qpt(torch.zeros((1, 2)), 0.25, 128), | |
| qpt(torch.zeros((1, 2)), 0.25, 128), | |
| ] | |
| ) | |
| # NOTE: NNAPI qadd supports broadcast, but PT does not. | |
| def test_qlinear(self): | |
| torch.manual_seed(29) | |
| weight = qpt(torch.randn(16, 32), 0.125, 0, torch.qint8) | |
| bias = torch.randn(16) | |
| mod = torch.ao.nn.quantized.Linear(32, 16) | |
| mod.set_weight_bias(weight, bias) | |
| inp = qpt(torch.randn(2, 32), 0.05, 130, torch.quint8) | |
| self.check(mod, inp) | |
| def test_seblock_mul(self): | |
| class MulModel(torch.nn.Module): | |
| def forward(self, lhs, rhs): | |
| return lhs * rhs | |
| self.check( | |
| MulModel(), | |
| [ | |
| nhwc(torch.randn(2, 3, 4, 4)), | |
| torch.randn(1, 3, 1, 1), | |
| ]) | |
| def test_multi_output(self): | |
| class MultiModel(torch.nn.Module): | |
| def forward(self, lhs, rhs) -> Tuple[torch.Tensor, torch.Tensor]: | |
| the_sum = lhs + rhs | |
| the_diff = lhs - rhs | |
| return the_sum, the_diff | |
| self.check(MultiModel(), [torch.tensor([1.0, 2.0]), torch.tensor([1.0, 3.0])]) | |
| if __name__ == '__main__': | |
| run_tests() |