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()
