def test_single_layer(self):
r"""Compare the result of quantizing single linear layer in
eager mode and graph mode
"""
# eager mode
annotated_linear_model = AnnotatedSingleLayerLinearModel()
linear_model = SingleLayerLinearModel()
# copy the weight from eager mode so that we can
# compare the result of the two quantized models later
linear_model.fc1.weight = torch.nn.Parameter(annotated_linear_model.fc1.module.weight.detach())
linear_model.fc1.bias = torch.nn.Parameter(annotated_linear_model.fc1.module.bias.detach())
model_eager = quantize(annotated_linear_model, test_only_eval_fn,
self.calib_data)
qconfig_dict = {
'': QConfig(
activation=default_observer,
weight=default_weight_observer)
}
model_script = quantize_script(
torch.jit.script(linear_model),
qconfig_dict,
test_only_eval_fn,
[self.calib_data],
inplace=False)
result_eager = model_eager(self.calib_data[0][0])
torch._C._jit_pass_quant_fusion(model_script._c._get_module('fc1')._get_method('forward').graph)
result_script = model_script._c._get_method('forward')(self.calib_data[0][0])
self.assertEqual(result_eager, result_script)
def quantize_and_run(model_eager, raw_model, img_data, do_eager=False):
qconfig_dict = {'': default_qconfig}
model_traced = torch.jit.trace(raw_model, img_data[0][0])
torch._C._jit_pass_inline(model_traced.graph)
model_quantized = quantize_script(model_traced,
qconfig_dict,
default_eval_fn, [img_data],
inplace=False)
result_traced = model_quantized(img_data[0][0])
torch._C._jit_pass_inline(model_quantized.graph)
print(model_quantized.graph)
if do_eager:
quantize_model(model_eager, img_data[0][0])
result_eager = model_eager(img_data[0][0])
np.allclose(result_traced.numpy(), result_eager.numpy())
def test_single_layer(self):
r"""Quantize SingleLayerLinearModel which has one Linear module, make sure it is swapped
to nnq.Linear which is the quantized version of the module
"""
# eager mode
model_eager = quantize(AnnotatedSingleLayerLinearModel(),
test_only_eval_fn, self.calib_data)
qconfig_dict = {
'':
QConfig(activation=default_observer,
weight=default_weight_observer)
}
model_script = quantize_script(
torch.jit.script(SingleLayerLinearModel()), qconfig_dict,
test_only_eval_fn, [self.calib_data])
result_eager = model_eager(self.calib_data[0][0])
result_script = model_script._c._get_method('forward')(
self.calib_data[0][0])
self.assertEqual(result_eager, result_script)