Fix (fx): fix fx quantize for conv->bn

Xilinx · Jul 18, 2023 · b61ad76 · b61ad76
1 parent ae7da18
commit b61ad76
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/brevitas/graph/quantize_impl.py b/src/brevitas/graph/quantize_impl.py
@@ -46,6 +46,8 @@
     nn.PixelUnshuffle,
     nn.Identity)
 
+BATCH_NORM = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)
+
 
 def inp_placeholder_handler(model, input_quantizer):
     """
@@ -184,6 +186,18 @@ def output_quant_handler(
             user_module = get_module(model, user.target)
             if hasattr(user_module, 'act_quant'):
                 output_quant = False
+            elif isinstance(user_module, BATCH_NORM):
+                # If the user is BatchNorm, check BN's users and potentially requentize at
+                # the output of BN
+                output_quant = False
+                output_quant_handler(
+                    model,
+                    user,
+                    rewriters,
+                    is_sign_preserving,
+                    quant_identity_map,
+                    quant_act_map,
+                    unsigned_act_tuple)
         if output_quant:
             if quant_module_name is None and quant_module is None:
                 if is_sign_preserving and are_inputs_unsigned(