Feat (tests): adding tests for example

Xilinx · Feb 2, 2024 · f3d6fb3 · f3d6fb3
1 parent d9ef191
commit f3d6fb3
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/tests/brevitas_examples/test_examples_import.py b/tests/brevitas_examples/test_examples_import.py
@@ -55,3 +55,19 @@ def test_super_resolution_float_and_quant_models_match(upscale_factor, num_chann
     float_model = float_espcn(upscale_factor, num_channels)
     quant_model = quant_espcn(upscale_factor, num_channels, weight_quant=weight_quant)
     quant_model.load_state_dict(float_model.state_dict())
+
+
+@pytest.mark.parametrize(
+    "weight_quant",
+    [
+        Int8WeightPerChannelFloat,
+        Int8AccumulatorAwareWeightQuant,
+        Int8AccumulatorAwareZeroCenterWeightQuant])
+def test_image_classification_float_and_quant_models_match(weight_quant):
+    import brevitas.config as config
+    from brevitas_examples.imagenet_classification.a2q.resnet import float_resnet18
+    from brevitas_examples.imagenet_classification.a2q.resnet import quant_resnet18
+    config.IGNORE_MISSING_KEYS = True
+    float_model = float_resnet18(num_classes=10)
+    quant_model = quant_resnet18(num_classes=10, weight_quant=weight_quant)
+    quant_model.load_state_dict(float_model.state_dict())