loosen batchnorm merging restrictions, fix ternary handling

Author: jmitrevs

hls4ml/model/optimizer/__init__.py

@@ -61,6 +61,8 @@
         'merge_linear_activation',
         'fuse_batch_normalization',
         'eliminate_linear_activation',
+        'qkeras_factorize_alpha',
+        'extract_ternary_threshold',
         # The ones above here need to be before infer_precision_types
         'infer_precision_types',
         'channels_last_converter',
@@ -70,8 +72,6 @@
         'fuse_bias_add',
         'expand_layer_group',
         'output_rounding_saturation_mode',
-        'qkeras_factorize_alpha',
-        'extract_ternary_threshold',
     ],
     requires=['parse_qonnx'],
 )  # TODO Maybe not all QKeras optmizers belong here?

hls4ml/model/optimizer/passes/batchnorm_opt.py

@@ -170,12 +170,12 @@ def match(self, node):
             s1 = node.weights['scale'].data_unquantized
             b1 = node.weights['bias'].data_unquantized
             scale_compatible = (
-                (prev_node.get_attr('scale_quantizer') is None and node.get_attr('scale_quantizer') is None)
+                (prev_node.get_attr('scale_quantizer') is None or node.get_attr('scale_quantizer') is None)
                 or (s0 == np.ones_like(s0)).all()
                 or (s1 == np.ones_like(s1)).all()
             )
             bias_compatible = (
-                (prev_node.get_attr('bias_quantizer') is None and node.get_attr('bias_quantizer') is None)
+                (prev_node.get_attr('bias_quantizer') is None or node.get_attr('bias_quantizer') is None)
                 or (b0 == np.zeros_like(b0)).all()
                 or (b1 == np.zeros_like(b1)).all()
             )
@@ -195,26 +195,24 @@ def transform(self, model, node):
         # if len(node_map[node.outputs[0]]) > 1:
         #     return False
 
-        # only merge if the types are integer or fixed
-        if (
-            not isinstance(prev_node.weights['scale'].type.precision, (IntegerPrecisionType, FixedPrecisionType))
-            or not isinstance(prev_node.weights['bias'].type.precision, (IntegerPrecisionType, FixedPrecisionType))
-            or not isinstance(node.weights['scale'].type.precision, (IntegerPrecisionType, FixedPrecisionType))
-            or not isinstance(node.weights['bias'].type.precision, (IntegerPrecisionType, FixedPrecisionType))
-        ):
-            return False
-
         s0 = prev_node.weights['scale'].data_unquantized
         b0 = prev_node.weights['bias'].data_unquantized
         s1 = node.weights['scale'].data_unquantized
         b1 = node.weights['bias'].data_unquantized
 
-        s_quantizer = (
-            node.get_attr('scale_quantizer') if (s0 == np.ones_like(s0)).all() else prev_node.get_attr('scale_quantizer')
-        )
-        b_quantizer = (
-            node.get_attr('bias_quantizer') if (b0 == np.zeros_like(b0)).all() else prev_node.get_attr('bias_quantizer')
-        )
+        if (s0 == np.ones_like(s0)).all():
+            s_quantizer = node.get_attr('scale_quantizer')
+        elif (s1 == np.ones_like(s1)).all():
+            s_quantizer = prev_node.get_attr('scale_quantizer')
+        else:
+            s_quantizer = None
+
+        if (b0 == np.ones_like(b0)).all():
+            b_quantizer = node.get_attr('bias_quantizer')
+        elif (b1 == np.ones_like(b1)).all():
+            b_quantizer = prev_node.get_attr('bias_quantizer')
+        else:
+            b_quantizer = None
 
         node.set_attr('scale_quantizer', s_quantizer)
         node.set_attr('bias_quantizer', b_quantizer)

hls4ml/model/optimizer/passes/qkeras.py

@@ -163,8 +163,16 @@ def transform(self, model, node):
         else:
             n_in = node.get_attr('n_out')
 
+        # the name of the new ApplyAlpha node
+        alpha_name = node.get_attr('name') + '_alpha'
+
+        # make the precision auto
+        alpha_precision = {'Precision': 'auto'}
+        model.config.set_name_config(alpha_name, alpha_precision)
+        model.config.parse_name_config(alpha_name, alpha_precision)
+
         attrs = {
-            'name': node.get_attr('name') + '_alpha',
+            'name': alpha_name,
             'class_name': 'Alpha',
             'inputs': node.outputs,
             'n_in': n_in,

test/pytest/test_qkeras.py

@@ -356,8 +356,10 @@ def test_relu_negative_slope(randX_1000_1, quantizer, backend, io_type):
     ],
 )
 def test_qactivation_kwarg(randX_100_10, activation_quantizer, weight_quantizer):
-    if activation_quantizer in ['binary', 'ternary']:
+    if activation_quantizer in ['binary']:
         name = 'bnbt_qdense_alpha'
+    elif activation_quantizer in ['ternary']:
+        name = 'bnbt_qdense_ternary_scale'
     else:
         name = f'qdense_{eval(activation_quantizer).__class__.__name__}'