Add support for parsing simple brevitas layers as part of pytorch models

hls4ml/converters/pytorch/convolution.py

@@ -1,8 +1,12 @@
-from hls4ml.converters.pytorch_to_hls import pytorch_handler
+import numpy as np
+
+from hls4ml.converters.pytorch_to_hls import addQuantizationParameters, convert_uaq_to_apfixed, pytorch_handler
 from hls4ml.converters.utils import compute_padding_1d_pytorch, compute_padding_2d_pytorch, parse_data_format
+from hls4ml.model.quantizers import BrevitasQuantizer
+from hls4ml.model.types import FixedPrecisionType
 
 
-@pytorch_handler('Conv1d')
+@pytorch_handler('Conv1d', 'QuantConv1d')
 def parse_conv1d_layer(operation, layer_name, input_names, input_shapes, node, class_object, data_reader, config):
     assert 'Conv1d' in operation
 
@@ -13,12 +17,50 @@ def parse_conv1d_layer(operation, layer_name, input_names, input_shapes, node, c
     layer['class_name'] = 'Conv1D'
     layer['data_format'] = 'channels_first'  # Pytorch default (can't change)
 
-    layer['weight_data'] = class_object.weight.data.numpy()
-    if class_object.bias is not None:
-        layer['bias_data'] = class_object.bias.data.numpy()
-    else:
-        layer['bias_data'] = None
+    if "Quant" in operation:
+        if class_object.weight_quant.is_quant_enabled:
+            width = int(class_object.quant_weight().bit_width)
+            scale = class_object.quant_weight().scale.detach().numpy()
+            mantissa, _ = np.frexp(scale)
+            # if scale is power of 2 we can simply use hls4ml FixedPrecisionType and directly
+            # use the already quantized tensor from brevitas
+            if mantissa == 0.5:
+                ap_fixed_params = convert_uaq_to_apfixed(width, float(class_object.quant_weight().scale))
+                layer['weight_data'] = class_object.quant_weight().detach().value.numpy()
+                layer['weight_quantizer'] = BrevitasQuantizer(
+                    width, FixedPrecisionType(width=width, integer=int(ap_fixed_params[1]), signed=True)
+                )
+            else:
+                raise Exception(
+                    '''Non-power of 2 quantization of weights not supported when injecting brevitas models.
+                    Please used QONNX instead.'''
+                )
+        else:
+            layer['weight_data'] = class_object.weight.data.numpy()
+
+        if class_object.bias_quant.is_quant_enabled:
+            width = int(class_object.quant_bias().bit_width)
+            ap_fixed_params = convert_uaq_to_apfixed(width, float(class_object.quant_bias().scale))
+            layer['bias_data'] = class_object.quant_bias().detach().value.numpy()
+            layer['bias_quantizer'] = BrevitasQuantizer(
+                width, FixedPrecisionType(width=width, integer=int(ap_fixed_params[1]), signed=True)
+            )
+        else:
+            if class_object.bias is not None:
+                layer['bias_data'] = class_object.bias.data.numpy()
+            else:
+                layer['bias_data'] = None
+        if class_object.input_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'input', act=True)
+        if class_object.output_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'output', act=True)
 
+    else:
+        layer['weight_data'] = class_object.weight.data.numpy()
+        if class_object.bias is not None:
+            layer['bias_data'] = class_object.bias.data.numpy()
+        else:
+            layer['bias_data'] = None
     # Input info
     (*_, layer['in_width'], layer['n_chan']) = parse_data_format(
         input_shapes[0], 'channels_first'
@@ -47,7 +89,7 @@ def parse_conv1d_layer(operation, layer_name, input_names, input_shapes, node, c
     return layer, output_shape
 
 
-@pytorch_handler('Conv2d')
+@pytorch_handler('Conv2d', 'QuantConv2d')
 def parse_conv2d_layer(operation, layer_name, input_names, input_shapes, node, class_object, data_reader, config):
     assert 'Conv2d' in operation
 
@@ -58,11 +100,52 @@ def parse_conv2d_layer(operation, layer_name, input_names, input_shapes, node, c
     layer['class_name'] = 'Conv2D'
     layer['data_format'] = 'channels_first'  # Pytorch default (can't change)
 
-    layer['weight_data'] = class_object.weight.data.numpy()
-    if class_object.bias is not None:
-        layer['bias_data'] = class_object.bias.data.numpy()
+    if "Quant" in operation:
+        if class_object.weight_quant.is_quant_enabled:
+            width = int(class_object.quant_weight().bit_width)
+            scale = class_object.quant_weight().scale.detach().numpy()
+            mantissa, _ = np.frexp(scale)
+            # if scale is power of 2 we can simply use hls4ml FixedPrecisionType and directly
+            # use the already quantized tensor from brevitas
+            if mantissa == 0.5:
+                ap_fixed_params = convert_uaq_to_apfixed(width, float(class_object.quant_weight().scale))
+                layer['weight_data'] = class_object.quant_weight().detach().value.numpy()
+                layer['weight_quantizer'] = BrevitasQuantizer(
+                    width, FixedPrecisionType(width=width, integer=int(ap_fixed_params[1]), signed=True)
+                )
+            else:
+                raise Exception(
+                    '''Non-power of 2 quantization of weights not supported when injecting brevitas models.
+                    Please used QONNX instead.'''
+                )
+                # layer = addQuantizationParameters(layer, class_object.quant_weight(), 'weight')
+                # layer['weight_data'] = class_object.quant_weight().detach().value.numpy()
+        else:
+            layer['weight_data'] = class_object.weight.data.numpy()
+
+        if class_object.bias_quant.is_quant_enabled:
+            width = int(class_object.quant_bias().bit_width)
+            ap_fixed_params = convert_uaq_to_apfixed(width, float(class_object.quant_bias().scale))
+            layer['bias_data'] = class_object.quant_bias().detach().value.numpy()
+            layer['bias_quantizer'] = BrevitasQuantizer(
+                width, FixedPrecisionType(width=width, integer=int(ap_fixed_params[1]), signed=True)
+            )
+        else:
+            if class_object.bias is not None:
+                layer['bias_data'] = class_object.bias.data.numpy()
+            else:
+                layer['bias_data'] = None
+        if class_object.input_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'input', act=True)
+        if class_object.output_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'output', act=True)
+
     else:
-        layer['bias_data'] = None
+        layer['weight_data'] = class_object.weight.data.numpy()
+        if class_object.bias is not None:
+            layer['bias_data'] = class_object.bias.data.numpy()
+        else:
+            layer['bias_data'] = None
 
     # Input info
     (*_, layer['in_height'], layer['in_width'], layer['n_chan']) = parse_data_format(

hls4ml/converters/pytorch/core.py

@@ -1,6 +1,8 @@
 import numpy as np
 
-from hls4ml.converters.pytorch_to_hls import pytorch_handler
+from hls4ml.converters.pytorch_to_hls import addQuantizationParameters, convert_uaq_to_apfixed, pytorch_handler
+from hls4ml.model.quantizers import BrevitasQuantizer
+from hls4ml.model.types import FixedPrecisionType
 
 
 @pytorch_handler('Constant')
@@ -20,7 +22,33 @@ def parse_constant_layer(operation, layer_name, node):
     return layer, output_shape
 
 
-@pytorch_handler('Linear')
+# A QuantIdentity layer does nothing but quantize its inputs. Insert `Quant` node to be processed by QONNX optimizers
+@pytorch_handler('QuantIdentity')
+def parse_quantidentity_layer(operation, layer_name, input_names, input_shapes, node, class_object, data_reader, config):
+    assert 'QuantIdentity' in operation
+
+    layer = {}
+    layer['inputs'] = input_names
+
+    layer['class_name'] = 'Quant'
+    layer['name'] = layer_name
+
+    if class_object.act_quant.is_quant_enabled:
+        layer['bitwidth'] = int(class_object.act_quant.bit_width())
+        layer['signed'] = class_object.act_quant.is_signed
+        layer['scale'] = np.full(np.array(input_shapes[0][1:]), class_object.act_quant.scale())
+        layer['zeropt'] = float(class_object.act_quant.zero_point())
+        layer['narrow'] = class_object.act_quant.is_narrow_range
+        layer['rounding_mode'] = class_object.act_quant.rounding_mode
+
+    else:
+        raise Exception('''QuantIdentify layer without act quant does nothing, please remove from model.''')
+    output_shape = input_shapes[0]
+
+    return layer, output_shape
+
+
+@pytorch_handler('Linear', 'QuantLinear')
 def parse_linear_layer(operation, layer_name, input_names, input_shapes, node, class_object, data_reader, config):
     assert 'Linear' in operation
 
@@ -36,6 +64,44 @@ def parse_linear_layer(operation, layer_name, input_names, input_shapes, node, c
     else:
         layer['bias_data'] = None
 
+    if "Quant" in operation:
+        if class_object.weight_quant.is_quant_enabled:
+            width = int(class_object.quant_weight().bit_width)
+            scale = class_object.quant_weight().scale.detach().numpy()
+            mantissa, _ = np.frexp(scale)
+            # if scale is power of 2 we can simply use hls4ml FixedPrecisionType and directly
+            # use the already quantized tensor from brevitas
+            if mantissa == 0.5:
+                ap_fixed_params = convert_uaq_to_apfixed(width, float(class_object.quant_weight().scale))
+                layer['weight_data'] = class_object.quant_weight().detach().value.numpy()
+                layer['weight_quantizer'] = BrevitasQuantizer(
+                    width, FixedPrecisionType(width=width, integer=int(ap_fixed_params[1]), signed=True)
+                )
+            else:
+                raise Exception(
+                    '''Non-power of 2 quantization of weights not supported when injecting brevitas models.
+                    Please used QONNX instead.'''
+                )
+        else:
+            layer['weight_data'] = class_object.weight.data.numpy()
+
+        if class_object.bias_quant.is_quant_enabled:
+            width = int(class_object.quant_bias().bit_width)
+            ap_fixed_params = convert_uaq_to_apfixed(width, float(class_object.quant_bias().scale))
+            layer['bias_data'] = class_object.quant_bias().detach().value.numpy()
+            layer['bias_quantizer'] = BrevitasQuantizer(
+                width, FixedPrecisionType(width=width, integer=int(ap_fixed_params[1]), signed=True)
+            )
+        else:
+            if class_object.bias is not None:
+                layer['bias_data'] = class_object.bias.data.numpy()
+            else:
+                layer['bias_data'] = None
+        if class_object.input_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'input', act=True)
+        if class_object.output_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'output', act=True)
+
     if class_object is not None:
         layer['n_in'] = class_object.in_features
         layer['n_out'] = class_object.out_features
@@ -54,7 +120,19 @@ def parse_linear_layer(operation, layer_name, input_names, input_shapes, node, c
     return layer, output_shape
 
 
-activation_layers = ['Softmax', 'ReLU', 'LeakyReLU', 'Threshold', 'ELU', 'PReLU', 'Sigmoid', 'Tanh']
+activation_layers = [
+    'Softmax',
+    'ReLU',
+    'LeakyReLU',
+    'Threshold',
+    'ELU',
+    'PReLU',
+    'Sigmoid',
+    'Tanh',
+    'QuantReLU',
+    'QuantSigmoid',
+    'QuantTanh',
+]
 
 
 @pytorch_handler(*activation_layers)
@@ -66,6 +144,12 @@ def parse_activation_layer(operation, layer_name, input_names, input_shapes, nod
     layer['name'] = layer_name
     layer['inputs'] = input_names
 
+    if "Quant" in operation:
+        layer['class_name'] = operation.split('Quant')[-1]
+        layer['activation'] = layer['class_name']
+        if class_object.act_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.act_quant, 'output', act=True)
+
     if node.op == 'call_module':
         if layer['class_name'] in ['ReLU', 'Sigmoid', 'Tanh']:
             layer['class_name'] = 'Activation'

hls4ml/converters/pytorch/merge.py

@@ -1,4 +1,4 @@
-from hls4ml.converters.pytorch_to_hls import pytorch_handler
+from hls4ml.converters.pytorch_to_hls import addQuantizationParameters, pytorch_handler
 
 concat_layers = ['cat', 'concat', 'concatenate']
 
@@ -28,7 +28,7 @@ def parse_concat_layer(operation, layer_name, input_names, input_shapes, node, c
     return layer, output_shape
 
 
-add_layers = ['add']
+add_layers = ['add', 'QuantEltwiseAdd']
 multiply_layers = ['mul', 'multiply']
 subtract_layers = ['sub', 'subtract']
 min_layers = ['fmin', 'minimum']
@@ -56,6 +56,12 @@ def parse_merge_layer(operation, layer_name, input_names, input_shapes, node, cl
 
     layer['inputs'] = input_names
 
+    if 'Quant' in operation:
+        if class_object.input_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'input', act=True)
+        if class_object.output_quant.is_quant_enabled:
+            layer = addQuantizationParameters(layer, class_object.input_quant, 'output', act=True, scale_up=True)
+
     output_shape = input_shapes[0][:]
 
     return layer, output_shape

hls4ml/converters/pytorch/pooling.py

@@ -1,7 +1,12 @@
 from hls4ml.converters.pytorch_to_hls import pytorch_handler
 from hls4ml.converters.utils import compute_padding_1d_pytorch, compute_padding_2d_pytorch, parse_data_format
 
-pooling_layers = ['MaxPool1d', 'MaxPool2d', 'AvgPool1d', 'AvgPool2d']
+pooling_layers = [
+    'MaxPool1d',
+    'MaxPool2d',
+    'AvgPool1d',
+    'AvgPool2d',
+]  # TODO add support for special quantized average pool layers
 
 
 @pytorch_handler(*pooling_layers)
@@ -10,9 +15,9 @@ def parse_pooling_layer(operation, layer_name, input_names, input_shapes, node,
 
     layer = {}
 
-    if operation == 'MaxPool1d':
+    if 'MaxPool1d' in operation:
         layer['class_name'] = 'MaxPooling1D'
-    if operation == 'MaxPool2d':
+    if 'MaxPool2d' in operation:
         layer['class_name'] = 'MaxPooling2D'
     if operation == 'AvgPool1d':
         layer['class_name'] = 'AveragePooling1D'