def test_compare_model_outputs_linear_static_fx(self): r"""Compare the output of linear layer in static quantized model and corresponding output of linear layer in float model """ qengine = torch.backends.quantized.engine qconfig = get_default_qconfig(qengine) qconfig_dict = {"": qconfig} float_model = SingleLayerLinearModel() float_model.eval() prepared_model = prepare_fx(float_model, qconfig_dict) prepared_float_model = copy.deepcopy(prepared_model) # Run calibration test_only_eval_fn(prepared_model, self.calib_data) q_model = convert_fx(prepared_model) linear_data = self.calib_data[0][0] expected_act_compare_dict_keys = {"x.stats", "fc1.stats"} self.compare_and_validate_model_outputs_results_fx( prepared_float_model, q_model, expected_act_compare_dict_keys, linear_data)
def test_compare_weights_linear_static_fx(self): r"""Compare the weights of float and static quantized linear layer""" def calibrate(model, calib_data): model.eval() with torch.no_grad(): for inp in calib_data: model(*inp) def compare_and_validate_results(float_model, q_model): weight_dict = compare_weights_fx(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) float_model = SingleLayerLinearModel() float_model.eval() qengine = torch.backends.quantized.engine qconfig = get_default_qconfig(qengine) qconfig_dict = {"": qconfig} prepared_model = prepare_fx(float_model, qconfig_dict) backup_prepared_model = copy.deepcopy(prepared_model) backup_prepared_model.eval() # Run calibration calibrate(prepared_model, self.calib_data) q_model = convert_fx(prepared_model) compare_and_validate_results(backup_prepared_model, q_model)
def test_compare_model_stub_linear_static_fx(self): r"""Compare the output of static quantized linear layer and its float shadow module""" qengine = torch.backends.quantized.engine qconfig = get_default_qconfig(qengine) qconfig_dict = {"": qconfig} float_model = SingleLayerLinearModel() float_model.eval() prepared_model = prepare_fx(float_model, qconfig_dict) prepared_float_model = copy.deepcopy(prepared_model) # Run calibration test_only_eval_fn(prepared_model, self.calib_data) q_model = convert_fx(prepared_model) linear_data = self.calib_data[0][0] module_swap_list = [nn.Linear] expected_ob_dict_keys = {"fc1.stats"} self.compare_and_validate_model_stub_results_fx( prepared_float_model, q_model, module_swap_list, expected_ob_dict_keys, linear_data, )
def __init__(self): super().__init__() self.obs1 = ModelReportObserver() self.mod1 = SingleLayerLinearModel() self.obs2 = ModelReportObserver() self.fc1 = torch.nn.Linear(5, 5).to(dtype=torch.float) self.relu = torch.nn.ReLU()
def test_save_load_state_dict_script(self): """ Tests that we can save and load state_dict for observers that are scripted in a quantized model. """ obs_list = [ MinMaxObserver, MovingAverageMinMaxObserver, HistogramObserver ] for obs in obs_list: model = SingleLayerLinearModel().eval() qconfig = QConfig(activation=default_observer, weight=obs) qconfig_dict = {'': qconfig} scripted = torch.jit.script(model) scripted = torch.quantization.prepare_jit(scripted, qconfig_dict) x = torch.rand(5, 5) scripted(x) obs_dict = torch.quantization.get_observer_state_dict(scripted) # Load stats scripted_2 = torch.jit.script(model) scripted_2 = torch.quantization.prepare_jit( scripted_2, qconfig_dict) torch.quantization.load_observer_state_dict(scripted_2, obs_dict) # Verify that state_dict matches exactly with original one. self.assertEqual(scripted.state_dict(), scripted_2.state_dict())
def test_quantization_doc_fx(self): path_from_pytorch = "docs/source/quantization.rst" unique_identifier = "FXPTQ API Example::" input_fp32 = SingleLayerLinearModel().get_example_inputs() global_inputs = {"UserModel": SingleLayerLinearModel, "input_fp32": input_fp32} code = self._get_code(path_from_pytorch, unique_identifier) self._test_code(code, global_inputs)
def test_remove_qconfig_observer_fx(self): r"""Remove activation_post_process node from fx prepred model""" float_model = SingleLayerLinearModel() float_model.eval() qengine = torch.backends.quantized.engine qconfig = get_default_qconfig(qengine) qconfig_dict = {"": qconfig} prepared_model = prepare_fx(float_model, qconfig_dict) prepared_float_model = copy.deepcopy(prepared_model) prepared_float_model.eval() model = remove_qconfig_observer_fx(prepared_float_model) modules = dict(model.named_modules()) for node in model.graph.nodes: if node.op == "call_module": self.assertFalse(is_activation_post_process(modules[node.target]))
def test_compare_shadow_activations_linear(self): test_cases = ( (SingleLayerLinearModel(), None), ( SingleLayerLinearDynamicModel(), {"object_type": [(nn.Linear, default_dynamic_qconfig)]}, ), ) for m, qconfig_dict in test_cases: m.eval() res = self._test_match_shadow_activations( m, (torch.randn(5, 5),), results_len=1, qconfig_dict=qconfig_dict)
def test_compare_weights_linear(self): test_cases = ( (SingleLayerLinearModel(), None), ( SingleLayerLinearDynamicModel(), {"object_type": [(nn.Linear, default_dynamic_qconfig)]}, ), ) for m, qconfig_dict in test_cases: m.eval() res = self._test_extract_weights( m, results_len=1, qconfig_dict=qconfig_dict)
def test_compare_weights_linear_static_fx(self): r"""Compare the weights of float and static quantized linear layer""" qengine = torch.backends.quantized.engine qconfig = get_default_qconfig(qengine) qconfig_dict = {"": qconfig} float_model = SingleLayerLinearModel() float_model.eval() prepared_model = prepare_fx(float_model, qconfig_dict) prepared_float_model = copy.deepcopy(prepared_model) prepared_float_model.eval() # Run calibration test_only_eval_fn(prepared_model, self.calib_data) q_model = convert_fx(prepared_model) expected_weight_dict_keys = {"fc1._packed_params._packed_params"} self.compare_and_validate_model_weights_results_fx( prepared_float_model, q_model, expected_weight_dict_keys)
def __init__(self): super().__init__() self.obs1 = ModelReportObserver() self.nested = TinyNestModule() self.fc1 = SingleLayerLinearModel() self.relu = torch.nn.ReLU()