| import torch | |
| import torch.nn as nn | |
| import torch.nn.quantized as nnq | |
| from torch.quantization import ( | |
| DeQuantStub, | |
| QuantStub, | |
| convert, | |
| default_qconfig, | |
| prepare, | |
| quantize, | |
| quantize_dynamic, | |
| ) | |
| from torch.quantization._numeric_suite import ( | |
| OutputLogger, | |
| Shadow, | |
| ShadowLogger, | |
| compare_model_outputs, | |
| compare_model_stub, | |
| compare_weights, | |
| ) | |
| from torch.testing._internal.common_quantization import ( | |
| AnnotatedConvBnReLUModel, | |
| AnnotatedConvModel, | |
| AnnotatedSingleLayerLinearModel, | |
| LSTMwithHiddenDynamicModel, | |
| QuantizationTestCase, | |
| SingleLayerLinearDynamicModel, | |
| test_only_eval_fn, | |
| ) | |
| from torch.testing._internal.common_quantized import override_qengines | |
| class SubModule(torch.nn.Module): | |
| def __init__(self): | |
| super(SubModule, self).__init__() | |
| self.mod1 = nn.Identity() | |
| self.mod2 = nn.ReLU() | |
| def forward(self, x): | |
| x = self.mod1(x) | |
| x = self.mod2(x) | |
| return x | |
| class ModelWithSubModules(torch.nn.Module): | |
| def __init__(self): | |
| super(ModelWithSubModules, self).__init__() | |
| self.qconfig = default_qconfig | |
| self.mod1 = SubModule() | |
| self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) | |
| self.quant = QuantStub() | |
| self.dequant = DeQuantStub() | |
| def forward(self, x): | |
| x = self.quant(x) | |
| x = self.mod1(x) | |
| x = self.conv(x) | |
| x = self.dequant(x) | |
| return x | |
| class ModelWithFunctionals(torch.nn.Module): | |
| def __init__(self): | |
| super(ModelWithFunctionals, self).__init__() | |
| self.mycat = nnq.FloatFunctional() | |
| self.myadd = nnq.FloatFunctional() | |
| self.mymul = nnq.FloatFunctional() | |
| self.myadd_relu = nnq.FloatFunctional() | |
| self.my_scalar_add = nnq.FloatFunctional() | |
| self.my_scalar_mul = nnq.FloatFunctional() | |
| self.quant = QuantStub() | |
| self.dequant = DeQuantStub() | |
| def forward(self, x): | |
| x = self.quant(x) | |
| x = self.mycat.cat([x, x, x]) | |
| x = self.myadd.add(x, x) | |
| x = self.mymul.mul(x, x) | |
| x = self.myadd_relu.add_relu(x, x) | |
| w = self.my_scalar_add.add_scalar(x, -0.5) | |
| w = self.my_scalar_mul.mul_scalar(w, 0.5) | |
| w = self.dequant(w) | |
| return w | |
| class TestEagerModeNumericSuite(QuantizationTestCase): | |
| @override_qengines | |
| def test_compare_weights_conv_static(self): | |
| r"""Compare the weights of float and static quantized conv layer""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model): | |
| weight_dict = compare_weights( | |
| float_model.state_dict(), q_model.state_dict() | |
| ) | |
| self.assertEqual(len(weight_dict), 1) | |
| for k, v in weight_dict.items(): | |
| self.assertTrue(v["float"].shape == v["quantized"].shape) | |
| model_list = [AnnotatedConvModel(qengine), AnnotatedConvBnReLUModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize(model, test_only_eval_fn, [self.img_data_2d]) | |
| compare_and_validate_results(model, q_model) | |
| @override_qengines | |
| def test_compare_weights_linear_static(self): | |
| r"""Compare the weights of float and static quantized linear layer""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model): | |
| weight_dict = compare_weights( | |
| float_model.state_dict(), q_model.state_dict() | |
| ) | |
| self.assertEqual(len(weight_dict), 1) | |
| for k, v in weight_dict.items(): | |
| self.assertTrue(v["float"].shape == v["quantized"].shape) | |
| model_list = [AnnotatedSingleLayerLinearModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize(model, test_only_eval_fn, [self.calib_data]) | |
| compare_and_validate_results(model, q_model) | |
| @override_qengines | |
| def test_compare_weights_linear_dynamic(self): | |
| r"""Compare the weights of float and dynamic quantized linear layer""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model): | |
| weight_dict = compare_weights( | |
| float_model.state_dict(), q_model.state_dict() | |
| ) | |
| self.assertEqual(len(weight_dict), 1) | |
| for k, v in weight_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| model_list = [SingleLayerLinearDynamicModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize_dynamic(model) | |
| compare_and_validate_results(model, q_model) | |
| @override_qengines | |
| def test_compare_weights_lstm_dynamic(self): | |
| r"""Compare the weights of float and dynamic quantized LSTM layer""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model): | |
| weight_dict = compare_weights( | |
| float_model.state_dict(), q_model.state_dict() | |
| ) | |
| self.assertEqual(len(weight_dict), 1) | |
| for k, v in weight_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| model_list = [LSTMwithHiddenDynamicModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize_dynamic(model) | |
| compare_and_validate_results(model, q_model) | |
| @override_qengines | |
| def test_compare_model_stub_conv_static(self): | |
| r"""Compare the output of static quantized conv layer and its float shadow module""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model, module_swap_list, data): | |
| ob_dict = compare_model_stub(float_model, q_model, module_swap_list, data) | |
| self.assertEqual(len(ob_dict), 1) | |
| for k, v in ob_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| model_list = [AnnotatedConvModel(qengine), AnnotatedConvBnReLUModel(qengine)] | |
| module_swap_list = [nn.Conv2d, nn.intrinsic.modules.fused.ConvReLU2d] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize(model, test_only_eval_fn, [self.img_data_2d]) | |
| compare_and_validate_results( | |
| model, q_model, module_swap_list, self.img_data_2d[0][0] | |
| ) | |
| @override_qengines | |
| def test_compare_model_stub_linear_static(self): | |
| r"""Compare the output of static quantized linear layer and its float shadow module""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model, module_swap_list, data): | |
| ob_dict = compare_model_stub(float_model, q_model, module_swap_list, data) | |
| self.assertEqual(len(ob_dict), 1) | |
| for k, v in ob_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| linear_data = self.calib_data[0][0] | |
| module_swap_list = [nn.Linear] | |
| model_list = [AnnotatedSingleLayerLinearModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize(model, test_only_eval_fn, [self.calib_data]) | |
| compare_and_validate_results(model, q_model, module_swap_list, linear_data) | |
| @override_qengines | |
| def test_compare_model_stub_submodule_static(self): | |
| r"""Compare the output of static quantized submodule and its float shadow module""" | |
| qengine = torch.backends.quantized.engine | |
| model = ModelWithSubModules().eval() | |
| q_model = quantize(model, test_only_eval_fn, [self.img_data_2d]) | |
| module_swap_list = [SubModule] | |
| ob_dict = compare_model_stub( | |
| model, q_model, module_swap_list, self.img_data_2d[0][0] | |
| ) | |
| self.assertTrue(isinstance(q_model.mod1, Shadow)) | |
| self.assertFalse(isinstance(q_model.conv, Shadow)) | |
| for k, v in ob_dict.items(): | |
| for i, val in enumerate(v["quantized"]): | |
| torch.testing.assert_allclose( | |
| v["float"][i], v["quantized"][i].dequantize() | |
| ) | |
| @override_qengines | |
| def test_compare_model_stub_functional_static(self): | |
| r"""Compare the output of static quantized functional layer and its float shadow module""" | |
| qengine = torch.backends.quantized.engine | |
| model = ModelWithFunctionals().eval() | |
| model.qconfig = torch.quantization.get_default_qconfig("fbgemm") | |
| q_model = prepare(model, inplace=False) | |
| q_model(self.img_data_2d[0][0]) | |
| q_model = convert(q_model) | |
| module_swap_list = [nnq.FloatFunctional] | |
| ob_dict = compare_model_stub( | |
| model, q_model, module_swap_list, self.img_data_2d[0][0] | |
| ) | |
| self.assertEqual(len(ob_dict), 6) | |
| self.assertTrue(isinstance(q_model.mycat, Shadow)) | |
| self.assertTrue(isinstance(q_model.myadd, Shadow)) | |
| self.assertTrue(isinstance(q_model.mymul, Shadow)) | |
| self.assertTrue(isinstance(q_model.myadd_relu, Shadow)) | |
| self.assertTrue(isinstance(q_model.my_scalar_add, Shadow)) | |
| self.assertTrue(isinstance(q_model.my_scalar_mul, Shadow)) | |
| for k, v in ob_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| @override_qengines | |
| def test_compare_model_stub_linear_dynamic(self): | |
| r"""Compare the output of dynamic quantized linear layer and its float shadow module""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model, module_swap_list, data): | |
| ob_dict = compare_model_stub(float_model, q_model, module_swap_list, data) | |
| self.assertEqual(len(ob_dict), 1) | |
| for k, v in ob_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| linear_data = self.calib_data[0][0] | |
| model_list = [SingleLayerLinearDynamicModel(qengine)] | |
| module_swap_list = [nn.Linear, nn.LSTM] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize_dynamic(model) | |
| compare_and_validate_results(model, q_model, module_swap_list, linear_data) | |
| @override_qengines | |
| def test_compare_model_stub_lstm_dynamic(self): | |
| r"""Compare the output of dynamic quantized LSTM layer and its float shadow module""" | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results( | |
| float_model, q_model, module_swap_list, input, hidden | |
| ): | |
| ob_dict = compare_model_stub( | |
| float_model, q_model, module_swap_list, input, hidden | |
| ) | |
| self.assertEqual(len(ob_dict), 1) | |
| for k, v in ob_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| lstm_input = torch.rand((1, 1, 2)) | |
| lstm_hidden = (torch.rand(1, 1, 2), torch.rand(1, 1, 2)) | |
| model_list = [LSTMwithHiddenDynamicModel(qengine)] | |
| module_swap_list = [nn.Linear, nn.LSTM] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize_dynamic(model) | |
| compare_and_validate_results( | |
| model, q_model, module_swap_list, lstm_input, lstm_hidden | |
| ) | |
| @override_qengines | |
| def test_compare_model_outputs_conv_static(self): | |
| r"""Compare the output of conv layer in stataic quantized model and corresponding | |
| output of conv layer in float model | |
| """ | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model, data): | |
| act_compare_dict = compare_model_outputs(float_model, q_model, data) | |
| expected_act_compare_dict_keys = {"conv.stats", "quant.stats"} | |
| self.assertTrue(act_compare_dict.keys() == expected_act_compare_dict_keys) | |
| for k, v in act_compare_dict.items(): | |
| self.assertTrue(v["float"][0].shape == v["quantized"][0].shape) | |
| model_list = [AnnotatedConvModel(qengine), AnnotatedConvBnReLUModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize(model, test_only_eval_fn, [self.img_data_2d]) | |
| compare_and_validate_results(model, q_model, self.img_data_2d[0][0]) | |
| @override_qengines | |
| def test_compare_model_outputs_linear_static(self): | |
| r"""Compare the output of linear layer in static quantized model and corresponding | |
| output of conv layer in float model | |
| """ | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model, data): | |
| act_compare_dict = compare_model_outputs(float_model, q_model, data) | |
| expected_act_compare_dict_keys = {"fc1.quant.stats", "fc1.module.stats"} | |
| self.assertTrue(act_compare_dict.keys() == expected_act_compare_dict_keys) | |
| for k, v in act_compare_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| linear_data = self.calib_data[0][0] | |
| model_list = [AnnotatedSingleLayerLinearModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize(model, test_only_eval_fn, [self.calib_data]) | |
| compare_and_validate_results(model, q_model, linear_data) | |
| @override_qengines | |
| def test_compare_model_outputs_functional_static(self): | |
| r"""Compare the output of functional layer in static quantized model and corresponding | |
| output of conv layer in float model | |
| """ | |
| qengine = torch.backends.quantized.engine | |
| model = ModelWithFunctionals().eval() | |
| model.qconfig = torch.quantization.get_default_qconfig("fbgemm") | |
| q_model = prepare(model, inplace=False) | |
| q_model(self.img_data_2d[0][0]) | |
| q_model = convert(q_model) | |
| act_compare_dict = compare_model_outputs(model, q_model, self.img_data_2d[0][0]) | |
| self.assertEqual(len(act_compare_dict), 7) | |
| expected_act_compare_dict_keys = { | |
| "mycat.stats", | |
| "myadd.stats", | |
| "mymul.stats", | |
| "myadd_relu.stats", | |
| "my_scalar_add.stats", | |
| "my_scalar_mul.stats", | |
| "quant.stats", | |
| } | |
| self.assertTrue(act_compare_dict.keys() == expected_act_compare_dict_keys) | |
| for k, v in act_compare_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| @override_qengines | |
| def test_compare_model_outputs_linear_dynamic(self): | |
| r"""Compare the output of linear layer in dynamic quantized model and corresponding | |
| output of conv layer in float model | |
| """ | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model, data): | |
| act_compare_dict = compare_model_outputs(float_model, q_model, data) | |
| expected_act_compare_dict_keys = {"fc1.stats"} | |
| self.assertTrue(act_compare_dict.keys() == expected_act_compare_dict_keys) | |
| for k, v in act_compare_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(v["float"][i].shape == v["quantized"][i].shape) | |
| linear_data = self.calib_data[0][0] | |
| model_list = [SingleLayerLinearDynamicModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize_dynamic(model) | |
| compare_and_validate_results(model, q_model, linear_data) | |
| @override_qengines | |
| def test_compare_model_outputs_lstm_dynamic(self): | |
| r"""Compare the output of LSTM layer in dynamic quantized model and corresponding | |
| output of conv layer in float model | |
| """ | |
| qengine = torch.backends.quantized.engine | |
| def compare_and_validate_results(float_model, q_model, input, hidden): | |
| act_compare_dict = compare_model_outputs( | |
| float_model, q_model, input, hidden | |
| ) | |
| expected_act_compare_dict_keys = {"lstm.stats"} | |
| self.assertTrue(act_compare_dict.keys() == expected_act_compare_dict_keys) | |
| for k, v in act_compare_dict.items(): | |
| self.assertTrue(len(v["float"]) == len(v["quantized"])) | |
| for i, val in enumerate(v["quantized"]): | |
| self.assertTrue(len(v["float"][i]) == len(v["quantized"][i])) | |
| if i == 0: | |
| self.assertTrue(v["float"][i][0].shape == v["quantized"][i][0].shape) | |
| else: | |
| self.assertTrue( | |
| v["float"][i][0].shape == v["quantized"][i][0].shape | |
| ) | |
| self.assertTrue( | |
| v["float"][i][1].shape == v["quantized"][i][1].shape | |
| ) | |
| lstm_input = torch.rand((1, 1, 2)) | |
| lstm_hidden = (torch.rand(1, 1, 2), torch.rand(1, 1, 2)) | |
| model_list = [LSTMwithHiddenDynamicModel(qengine)] | |
| for model in model_list: | |
| model.eval() | |
| if hasattr(model, "fuse_model"): | |
| model.fuse_model() | |
| q_model = quantize_dynamic(model) | |
| compare_and_validate_results(model, q_model, lstm_input, lstm_hidden) | |
| @override_qengines | |
| def test_output_logger(self): | |
| r"""Compare output from OutputLogger with the expected results""" | |
| x = torch.rand(2, 2) | |
| y = torch.rand(2, 1) | |
| l = [] | |
| l.append(x) | |
| l.append(y) | |
| logger = OutputLogger() | |
| logger.forward(x) | |
| logger.forward(y) | |
| self.assertEqual(l, logger.stats["tensor_val"]) | |
| @override_qengines | |
| def test_shadow_logger(self): | |
| r"""Compare output from ShawdowLogger with the expected results""" | |
| a_float = torch.rand(2, 2) | |
| a_quantized = torch.rand(2, 2) | |
| b_float = torch.rand(3, 2, 2) | |
| b_quantized = torch.rand(3, 2, 2) | |
| logger = ShadowLogger() | |
| logger.forward(a_float, a_quantized) | |
| logger.forward(b_float, b_quantized) | |
| self.assertEqual(len(logger.stats["float"]), 2) | |
| self.assertEqual(len(logger.stats["quantized"]), 2) |