def test_compare_model_outputs_conv_static_fx(self):
r"""Compare the output of conv layer in static quantized model and corresponding
output of conv layer in float model
"""
qengine = torch.backends.quantized.engine
qconfig = get_default_qconfig(qengine)
qconfig_dict = {"": qconfig}
model_list = [ConvModel(), ConvBnReLUModel()]
for float_model in model_list:
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.img_data_2d)
q_model = convert_fx(prepared_model)
expected_act_compare_dict_keys = {"x.stats", "conv.stats"}
self.compare_and_validate_model_outputs_results_fx(
prepared_float_model,
q_model,
expected_act_compare_dict_keys,
self.img_data_2d[0][0],
)
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_model_stub_conv_static_fx(self):
r"""Compare the output of static quantized conv layer and its float shadow module"""
qengine = torch.backends.quantized.engine
qconfig = get_default_qconfig(qengine)
qconfig_dict = {"": qconfig}
model_list = [ConvModel(), ConvBnReLUModel()]
for float_model in model_list:
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.img_data_2d)
q_model = convert_fx(prepared_model)
module_swap_list = [nn.Conv2d, nni.modules.fused.ConvReLU2d]
expected_ob_dict_keys = {"conv.stats"}
self.compare_and_validate_model_stub_results_fx(
prepared_float_model,
q_model,
module_swap_list,
expected_ob_dict_keys,
self.img_data_2d[0][0],
)
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 checkQuantized(model):
self.assertEqual(type(model.conv), nnq.Conv2d)
self.assertEqual(type(model.fc1), nnq.Linear)
self.assertEqual(type(model.fc2), nnq.Linear)
test_only_eval_fn(model, self.img_data_2d)
self.checkScriptable(model, self.img_data_2d)
self.checkNoQconfig(model)
def test_dynamic_qat_linear(self):
for qengine in supported_qengines:
with override_quantized_engine(qengine):
# Dynamic QAT without memoryless observers should fail
with self.assertRaisesRegex(
ValueError,
"Dynamic QAT requires a memoryless observer." +
"This means a MovingAverage observer with averaging constant equal to 1"
):
model = ManualLinearDynamicQATModel(default_qat_qconfig)
model = prepare_qat(
model, mapping={torch.nn.Linear: nnqatd.Linear})
model = ManualLinearDynamicQATModel()
model = prepare_qat(model,
mapping={torch.nn.Linear: nnqatd.Linear})
self.assertEqual(type(model.fc1), nnqatd.Linear)
self.assertEqual(type(model.fc2), nnqatd.Linear)
self.checkObservers(model)
test_only_train_fn(model, self.train_data)
model = convert(model, mapping={nnqatd.Linear: nnqd.Linear})
self.assertEqual(type(model.fc1), nnqd.Linear)
self.assertEqual(type(model.fc2), nnqd.Linear)
test_only_eval_fn(model, self.calib_data)
self.checkScriptable(model, self.calib_data)
self.checkNoQconfig(model)
def checkQuantized(model):
# Make sure EmbeddingBag is now a quantized EmbeddingBag.
self.assertTrue(type(model.emb), nn.quantized.EmbeddingBag)
# Also test that Linear has been quantized.
self.assertTrue(type(model.linear), nnq.Linear)
test_only_eval_fn(model, self.embed_data)
self.checkScriptable(model, self.embed_data)
self.checkNoQconfig(model)
def checkQuantized(model):
# make sure Embedding is now a QuantizedEmbedding
self.assertEqual(type(model.emb), nn.quantized.Embedding)
# make sure Linear is now a QuantizedLinear
self.assertEqual(type(model.linear), nn.quantized.Linear)
test_only_eval_fn(model, self.embed_data)
self.checkScriptable(model, self.embed_data)
self.checkNoQconfig(model)
def checkQuantized(model):
self.assertEqual(type(model.conv1), nniq.ConvReLU2d)
self.assertEqual(type(model.bn1), nn.Identity)
self.assertEqual(type(model.relu1), nn.Identity)
self.assertEqual(type(model.sub1.conv), nnq.Conv2d)
self.assertEqual(type(model.sub1.bn), nn.Identity)
self.assertEqual(type(model.sub2.conv), nn.Conv2d)
self.assertEqual(type(model.sub2.relu), nn.ReLU)
test_only_eval_fn(model, self.img_data_1d)
self.checkNoQconfig(model)
def test_forward_hooks_preserved(self):
r"""Test case that checks whether forward pre hooks of the first module and
post forward hooks of the last module in modules list passed to fusion function preserved.
(e.g. before fusion: [nn.Conv2d (with pre forward hooks), nn.BatchNorm2d, nn.ReLU (with post forward hooks)]
after fusion: [nni.ConvBnReLU2d (with pre and post hooks), nn.Identity, nn.Identity])
"""
model = ModelForFusion(default_qat_qconfig).train()
counter = {
'pre_forwards': 0,
'forwards': 0,
}
fused = False
def fw_pre_hook(fused_module_class, h_module, input):
if fused:
self.assertEqual(
type(h_module), fused_module_class,
"After fusion owner of the first module's forward pre hook is not a fused module"
)
counter['pre_forwards'] += 1
def fw_hook(fused_module_class, h_module, input, output):
if fused:
self.assertEqual(
type(h_module), fused_module_class,
"After fusion owner of the last module's forward hook is not a fused module"
)
counter['forwards'] += 1
# Registering two pre and two post forward hooks, thus expecting counter increment by two each inference
model.conv1.register_forward_pre_hook(
lambda *args: fw_pre_hook(nni.ConvBnReLU2d, *args))
model.sub1.conv.register_forward_pre_hook(
lambda *args: fw_pre_hook(nni.ConvBn2d, *args))
model.relu1.register_forward_hook(
lambda *args: fw_hook(nni.ConvBnReLU2d, *args))
model.sub1.bn.register_forward_hook(
lambda *args: fw_hook(nni.ConvBn2d, *args))
test_only_eval_fn(model, self.img_data_1d)
self.assertEqual(counter['pre_forwards'], 2 * len(self.img_data_1d))
self.assertEqual(counter['forwards'], 2 * len(self.img_data_1d))
model = fuse_modules(model, ['conv1', 'bn1', 'relu1'])
model = fuse_modules(model, ['sub1.conv', 'sub1.bn'])
fused = True
before_fusion_pre_count = counter['pre_forwards']
before_fusion_post_count = counter['forwards']
test_only_eval_fn(model, self.img_data_1d)
self.assertEqual(counter['pre_forwards'] - before_fusion_pre_count,
2 * len(self.img_data_1d))
self.assertEqual(counter['forwards'] - before_fusion_post_count,
2 * len(self.img_data_1d))
def test_eval_only_fake_quant(self):
r"""Using FakeQuant in evaluation only mode,
this is useful for estimating accuracy loss when we quantize the
network
"""
for qengine in supported_qengines:
with override_quantized_engine(qengine):
model = ManualLinearQATModel(qengine)
model = prepare_qat(model)
self.checkObservers(model)
model.eval()
test_only_eval_fn(model, self.calib_data)
def test_record_observer(self):
for qengine in supported_qengines:
with override_quantized_engine(qengine):
model = AnnotatedSingleLayerLinearModel()
model.qconfig = default_debug_qconfig
model = prepare(model)
# run the evaluation and dump all tensors
test_only_eval_fn(model, self.calib_data)
test_only_eval_fn(model, self.calib_data)
observer_dict = {}
get_observer_dict(model, observer_dict)
self.assertTrue('fc1.module.activation_post_process' in observer_dict.keys(),
'observer is not recorded in the dict')
self.assertEqual(len(observer_dict['fc1.module.activation_post_process'].get_tensor_value()),
2 * len(self.calib_data))
self.assertEqual(observer_dict['fc1.module.activation_post_process'].get_tensor_value()[0],
model(self.calib_data[0][0]))
def test_compare_weights_conv_static_fx(self):
r"""Compare the weights of float and static quantized conv layer"""
qengine = torch.backends.quantized.engine
qconfig = get_default_qconfig(qengine)
qconfig_dict = {"": qconfig}
model_list = [ConvModel(), ConvBnModel(), ConvBnReLUModel()]
for float_model in model_list:
float_model.eval()
fused = fuse_fx(float_model)
prepared_model = prepare_fx(float_model, qconfig_dict)
# Run calibration
test_only_eval_fn(prepared_model, self.img_data_2d)
q_model = convert_fx(prepared_model)
expected_weight_dict_keys = {"conv.weight"}
self.compare_and_validate_model_weights_results_fx(
fused, q_model, expected_weight_dict_keys)
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 test_fuse_module_eval(self):
model = ModelForFusion(default_qconfig)
model.eval()
model = fuse_modules(
model,
[['conv3', 'bn3', 'relu4'], ['conv1', 'bn1', 'relu1'],
['conv2', 'relu2'], ['bn2', 'relu3'], ['sub1.conv', 'sub1.bn']])
self.assertEqual(
type(model.conv1),
nni.ConvReLU2d,
msg="Fused Conv + BN + Relu first layer (BN is folded)")
self.assertEqual(type(model.conv1[0]),
nn.Conv2d,
msg="Fused Conv + BN + Relu (Conv + folded BN only)")
self.assertEqual(type(model.conv1[1]),
nn.ReLU,
msg="Fused Conv + BN + Relu second layer (Relu only)")
self.assertEqual(
type(model.bn1),
nn.Identity,
msg="Fused Conv + BN + Relu second layer (Skipped BN)")
self.assertEqual(
type(model.relu1),
nn.Identity,
msg="Fused Conv + BN + Relu second layer (Skipped Relu)")
self.assertEqual(
type(model.conv2),
nni.ConvReLU3d,
msg="Fused Conv + BN + Relu first layer (BN is folded)")
self.assertEqual(type(model.bn2),
nni.BNReLU3d,
msg="Fused BN + Relu first layer (Relu is folded))")
self.assertEqual(type(model.relu3),
nn.Identity,
msg="Fused BN + Relu second layer (Skipped Relu)")
self.assertEqual(type(model.conv2[0]),
nn.Conv3d,
msg="Fused Conv + BN + Relu (Conv + folded BN only)")
self.assertEqual(type(model.conv2[1]),
nn.ReLU,
msg="Fused Conv + BN + Relu second layer (Relu only)")
self.assertEqual(
type(model.relu2),
nn.Identity,
msg="Fused Conv + BN + Relu second layer (Skipped Relu)")
self.assertEqual(type(model.conv3),
nni.ConvReLU1d,
msg="Fused Conv + Relu for Conv1d (folded BN)")
self.assertEqual(type(model.conv3[0]),
nn.Conv1d,
msg="Fused Conv + Relu for Conv1d ")
self.assertEqual(type(model.conv3[1]),
nn.ReLU,
msg="Fused Conv + Relu for Conv1d")
self.assertEqual(type(model.bn3),
nn.Identity,
msg="Fused Conv + BN + Relu for Conv1d (Skipped BN)")
self.assertEqual(type(model.sub1.conv),
nn.Conv2d,
msg="Fused submodule Conv + folded BN")
self.assertEqual(type(model.sub1.bn),
nn.Identity,
msg="Fused submodule (skipped BN)")
self.assertEqual(type(model.sub2.conv),
nn.Conv2d,
msg="Non-fused submodule Conv")
self.assertEqual(type(model.sub2.relu),
torch.nn.ReLU,
msg="Non-fused submodule ReLU")
model = prepare(model)
self.checkObservers(model)
test_only_eval_fn(model, self.img_data_1d)
model = convert(model)
def checkQuantized(model):
self.assertEqual(type(model.conv3), nniq.ConvReLU1d)
self.assertEqual(type(model.conv1), nniq.ConvReLU2d)
self.assertEqual(type(model.bn1), nn.Identity)
self.assertEqual(type(model.relu1), nn.Identity)
self.assertEqual(type(model.sub1.conv), nnq.Conv2d)
self.assertEqual(type(model.sub1.bn), nn.Identity)
self.assertEqual(type(model.sub2.conv), nn.Conv2d)
self.assertEqual(type(model.sub2.relu), nn.ReLU)
self.assertEqual(type(model.bn2), nniq.BNReLU3d)
test_only_eval_fn(model, self.img_data_1d)
self.checkNoQconfig(model)
checkQuantized(model)
model = ModelForFusion(default_qconfig).eval()
model = fuse_modules(
model,
[['conv1', 'bn1', 'relu1'], ['conv2', 'relu2'], ['bn2', 'relu3'],
['sub1.conv', 'sub1.bn'], ['conv3', 'bn3', 'relu4']])
model = quantize(model, test_only_eval_fn, [self.img_data_1d])
checkQuantized(model)
def checkQuantized(model):
self.assertEqual(type(model.fc1), nnq.Linear)
self.assertEqual(type(model.dropout), nnq.Dropout)
test_only_eval_fn(model, self.calib_data)
self.checkScriptable(model, self.calib_data)
self.checkNoQconfig(model)
