def test_compare_model_stub(self):
r"""Compare the output of quantized conv layer and its float shadow module
"""
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, ShadowLogger)
self.assertEqual(len(ob_dict), 1)
for k, v in ob_dict.items():
self.assertTrue(v["float"].shape == v["quantized"].shape)
for qengine in supported_qengines:
with override_quantized_engine(qengine):
model_list = [
AnnotatedConvModel(qengine),
AnnotatedConvBnReLUModel(qengine),
]
data = self.img_data[0][0]
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, default_eval_fn, self.img_data)
compare_and_validate_results(model, q_model,
module_swap_list, data)
# Test adding stub to sub module
model = ModelWithSubModules().eval()
q_model = quantize(model, default_eval_fn, self.img_data)
module_swap_list = [SubModule]
ob_dict = compare_model_stub(model, q_model, module_swap_list,
data, ShadowLogger)
self.assertTrue(isinstance(q_model.mod1, Shadow))
self.assertFalse(isinstance(q_model.conv, Shadow))
for k, v in ob_dict.items():
torch.testing.assert_allclose(v["float"],
v["quantized"].dequantize())
# Test adding stub to functionals
model = ModelWithFunctionals().eval()
model.qconfig = torch.quantization.get_default_qconfig(
"fbgemm")
q_model = prepare(model, inplace=False)
q_model(data)
q_model = convert(q_model)
module_swap_list = [nnq.FloatFunctional]
ob_dict = compare_model_stub(model, q_model, module_swap_list,
data, ShadowLogger)
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(v["float"].shape == v["quantized"].shape)
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, ShadowLogger)
self.assertEqual(len(ob_dict), 1)
for k, v in ob_dict.items():
self.assertTrue(v["float"].shape == v["quantized"].shape)
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)
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():
torch.testing.assert_allclose(v["float"],
v["quantized"].dequantize())
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(v["float"].shape == v["quantized"].shape)
float_model.eval()
float_model.fuse_model()
float_model.qconfig = torch.quantization.default_qconfig
img_data = [(torch.rand(2, 3, 10, 10, dtype=torch.float), torch.randint(0, 1, (2,), dtype=torch.long)) for _ in range(2)]
qmodel = quantize(float_model, default_eval_fn, [img_data], inplace=False)
##############################################################################
# In the following example we call ``compare_model_stub()`` from PyTorch Numeric Suite to compare ``QuantizableBasicBlock`` module with its float point equivalent. This API returns a dict with key corresponding to module names and each entry being a dictionary with two keys 'float' and 'quantized', containing the output tensors of quantized and its matching float shadow module.
data = img_data[0][0]
module_swap_list = [torchvision.models.quantization.resnet.QuantizableBasicBlock]
# Takes in floating point and quantized model as well as input data, and returns a dict with key
# corresponding to module names and each entry being a dictionary with two keys 'float' and
# 'quantized', containing the output tensors of quantized module and its matching floating point shadow module.
ob_dict = ns.compare_model_stub(float_model, qmodel, module_swap_list, data)
print('keys of ob_dict:')
print(ob_dict.keys())
print("\nkeys of ob_dict entry for layer1.0's output:")
print(ob_dict['layer1.0.stats'].keys())
print(ob_dict['layer1.0.stats']['float'][0].shape)
print(ob_dict['layer1.0.stats']['quantized'][0].shape)
##############################################################################
# This dict can be then used to compare and compute the module level quantization error.
for key in ob_dict:
print(key, compute_error(ob_dict[key]['float'][0], ob_dict[key]['quantized'][0].dequantize()))
