Correctly expand dims of pointwise layer (#715)

vloncar · web-flow · commit f0bcd4f02451 · 2023-03-06T13:17:34.000-08:00
* Correctly expand dims of pointwise layer

* Style rules applied
diff --git a/hls4ml/backends/quartus/passes/pointwise.py b/hls4ml/backends/quartus/passes/pointwise.py
@@ -1,39 +1,56 @@
-import numpy as np
 from copy import copy
 
-from hls4ml.model.optimizer import OptimizerPass
-from hls4ml.model.layers import register_layer
+import numpy as np
+
 from hls4ml.backends.fpga.fpga_layers import PointwiseConv1D, PointwiseConv2D
-from hls4ml.backends.quartus.passes.convolution_templates import  Conv1DConfigTemplate, Conv1DFunctionTemplate, Conv2DConfigTemplate, Conv2DFunctionTemplate, conv1d_config_template, conv2d_config_template, conv_mult_config_template
+from hls4ml.backends.quartus.passes.convolution_templates import (
+    Conv1DConfigTemplate,
+    Conv1DFunctionTemplate,
+    Conv2DConfigTemplate,
+    Conv2DFunctionTemplate,
+    conv1d_config_template,
+    conv2d_config_template,
+    conv_mult_config_template,
+)
+from hls4ml.model.layers import register_layer
+from hls4ml.model.optimizer import OptimizerPass
 
 '''
 Custom hls4ml layer implementation for 1x1 Conv filters using im2col
 Allows lower latency andresource usage, due to less loop invocations
 '''
 
-pointwise_conv1d_function_template = 'nnet::pointwise_conv_1d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
-pointwise_conv2d_function_template = 'nnet::pointwise_conv_2d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
+pointwise_conv1d_function_template = (
+    'nnet::pointwise_conv_1d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
+)
+pointwise_conv2d_function_template = (
+    'nnet::pointwise_conv_2d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
+)
 
 sepconv1d_include_list = ['nnet_utils/nnet_conv1d.h']
 sepconv2d_include_list = ['nnet_utils/nnet_conv2d.h']
 
+
 class PointwiseConv1DConfigTemplate(Conv1DConfigTemplate):
     def __init__(self):
         super(Conv1DConfigTemplate, self).__init__(PointwiseConv1D)
         self.template = conv1d_config_template
         self.mult_template = conv_mult_config_template
 
+
 class PointwiseConv1DFunctionTemplate(Conv1DFunctionTemplate):
     def __init__(self):
         super(Conv1DFunctionTemplate, self).__init__(PointwiseConv1D, include_header=sepconv1d_include_list)
         self.template = pointwise_conv1d_function_template
 
+
 class PointwiseConv2DConfigTemplate(Conv2DConfigTemplate):
     def __init__(self):
         super(Conv2DConfigTemplate, self).__init__(PointwiseConv2D)
         self.template = conv2d_config_template
         self.mult_template = conv_mult_config_template
 
+
 class PointwiseConv2DFunctionTemplate(Conv2DFunctionTemplate):
     def __init__(self):
         super(Conv2DFunctionTemplate, self).__init__(PointwiseConv2D, include_header=sepconv2d_include_list)
@@ -54,19 +71,25 @@ def register_pointwise(backend):
     backend.register_template(PointwiseConv2DConfigTemplate)
     backend.register_template(PointwiseConv2DFunctionTemplate)
 
+
 class OptimizePointwiseConv(OptimizerPass):
     def match(self, node):
-        return node.class_name in ('Conv1D', 'Conv2D') and \
-            node.get_attr('filt_height', 1) == 1 and \
-            node.get_attr('filt_width') == 1 and \
-            node.model.config.get_config_value('IOType') == 'io_parallel'
+        return (
+            node.class_name in ('Conv1D', 'Conv2D')
+            and node.get_attr('filt_height', 1) == 1
+            and node.get_attr('filt_width') == 1
+            and node.model.config.get_config_value('IOType') == 'io_parallel'
+        )
 
     def transform(self, model, node):
-        dim = node.__class__.__name__[-2:] # '1D' or '2D'
-        pw_node = model.make_node('PointwiseConv' + dim, node.name, copy(node.attributes), node.inputs.copy(), outputs=node.outputs.copy())
-        if len(node.weights['weight'].data.shape) == 2: # This can happen if we assign weights of Dense layer to 1x1 Conv2D
-            pw_node.weights['weight'].data = np.expand_dims(node.weights['weight'].data, axis=(0,1))
+        dim = node.__class__.__name__[-2:]  # '1D' or '2D'
+        pw_node = model.make_node(
+            'PointwiseConv' + dim, node.name, copy(node.attributes), node.inputs.copy(), outputs=node.outputs.copy()
+        )
+        if len(node.weights['weight'].data.shape) == 2:  # This can happen if we assign weights of Dense layer to 1x1 Conv2D
+            expand_axis = tuple(range(int(dim[0])))
+            pw_node.weights['weight'].data = np.expand_dims(node.weights['weight'].data, axis=expand_axis)
         pw_node.weights['bias'].data = node.weights['bias'].data
         model.replace_node(node, pw_node)
-        
+
         return True
diff --git a/hls4ml/backends/vivado/passes/pointwise.py b/hls4ml/backends/vivado/passes/pointwise.py
@@ -1,34 +1,51 @@
-import numpy as np
 from copy import copy
 
-from hls4ml.model.optimizer import OptimizerPass
-from hls4ml.model.layers import register_layer
+import numpy as np
+
 from hls4ml.backends.fpga.fpga_layers import PointwiseConv1D, PointwiseConv2D
-from hls4ml.backends.vivado.passes.convolution_templates import Conv1DConfigTemplate, Conv1DFunctionTemplate, Conv2DConfigTemplate, Conv2DFunctionTemplate, conv1d_config_template, conv2d_config_template, conv_mult_config_template
+from hls4ml.backends.vivado.passes.convolution_templates import (
+    Conv1DConfigTemplate,
+    Conv1DFunctionTemplate,
+    Conv2DConfigTemplate,
+    Conv2DFunctionTemplate,
+    conv1d_config_template,
+    conv2d_config_template,
+    conv_mult_config_template,
+)
+from hls4ml.model.layers import register_layer
+from hls4ml.model.optimizer import OptimizerPass
 
-pointwise_conv1d_function_template = 'nnet::pointwise_conv_1d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
-pointwise_conv2d_function_template = 'nnet::pointwise_conv_2d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
+pointwise_conv1d_function_template = (
+    'nnet::pointwise_conv_1d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
+)
+pointwise_conv2d_function_template = (
+    'nnet::pointwise_conv_2d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
+)
 
 sepconv1d_include_list = ['nnet_utils/nnet_conv1d.h', 'nnet_utils/nnet_sepconv1d_stream.h']
 sepconv2d_include_list = ['nnet_utils/nnet_conv2d.h', 'nnet_utils/nnet_sepconv2d_stream.h']
 
+
 class PointwiseConv1DConfigTemplate(Conv1DConfigTemplate):
     def __init__(self):
         super(Conv1DConfigTemplate, self).__init__(PointwiseConv1D)
         self.template = conv1d_config_template
         self.mult_template = conv_mult_config_template
 
+
 class PointwiseConv1DFunctionTemplate(Conv1DFunctionTemplate):
     def __init__(self):
         super(Conv1DFunctionTemplate, self).__init__(PointwiseConv1D, include_header=sepconv1d_include_list)
         self.template = pointwise_conv1d_function_template
 
+
 class PointwiseConv2DConfigTemplate(Conv2DConfigTemplate):
     def __init__(self):
         super(Conv2DConfigTemplate, self).__init__(PointwiseConv2D)
         self.template = conv2d_config_template
         self.mult_template = conv_mult_config_template
 
+
 class PointwiseConv2DFunctionTemplate(Conv2DFunctionTemplate):
     def __init__(self):
         super(Conv2DFunctionTemplate, self).__init__(PointwiseConv2D, include_header=sepconv2d_include_list)
@@ -49,18 +66,22 @@ def register_pointwise(backend):
     backend.register_template(PointwiseConv2DConfigTemplate)
     backend.register_template(PointwiseConv2DFunctionTemplate)
 
+
 class OptimizePointwiseConv(OptimizerPass):
     def match(self, node):
-        return node.class_name in ('Conv1D', 'Conv2D') and \
-            node.get_attr('filt_height', 1) == 1 and \
-            node.get_attr('filt_width') == 1
+        return (
+            node.class_name in ('Conv1D', 'Conv2D')
+            and node.get_attr('filt_height', 1) == 1
+            and node.get_attr('filt_width') == 1
+        )
 
     def transform(self, model, node):
-        dim = node.__class__.__name__[-2:] # '1D' or '2D'
+        dim = node.__class__.__name__[-2:]  # '1D' or '2D'
         pw_node = model.make_node('PointwiseConv' + dim, node.name, copy(node.attributes), node.inputs.copy())
-        if len(node.weights['weight'].data.shape) == 2: # This can happen if we assign weights of Dense layer to 1x1 Conv2D
-            pw_node.weights['weight'].data = np.expand_dims(node.weights['weight'].data, axis=(0,1))
+        if len(node.weights['weight'].data.shape) == 2:  # This can happen if we assign weights of Dense layer to 1x1 Conv2D
+            expand_axis = tuple(range(int(dim[0])))
+            pw_node.weights['weight'].data = np.expand_dims(node.weights['weight'].data, axis=expand_axis)
         pw_node.weights['bias'].data = node.weights['bias'].data
         model.replace_node(node, pw_node)
-        
-        return True
+
+        return True
diff --git a/hls4ml/model/optimizer/passes/multi_dense.py b/hls4ml/model/optimizer/passes/multi_dense.py
@@ -1,17 +1,21 @@
-from hls4ml.model.optimizer import OptimizerPass
-from hls4ml.model.layers import Dense
 import numpy as np
 
+from hls4ml.model.layers import Dense
+from hls4ml.model.optimizer import OptimizerPass
+
+
 class ReplaceMultidimensionalDenseWithConv(OptimizerPass):
     def match(self, node):
-        return isinstance(node, Dense) and \
-            len(node.get_input_variable().shape) - sum(d==1 for d in node.get_input_variable().shape) > 1 
-            # The above sum checks for the number of dimensions in the Dense with size 1
-            # The subtraction allows the check to only count the number of dimensions with non-1 size
-            # For example, this prevents matching for a Dense layer with shape (1,N)
+        return (
+            isinstance(node, Dense)
+            and len(node.get_input_variable().shape) - sum(d == 1 for d in node.get_input_variable().shape) > 1
+        )
+        # The above sum checks for the number of dimensions in the Dense with size 1
+        # The subtraction allows the check to only count the number of dimensions with non-1 size
+        # For example, this prevents matching for a Dense layer with shape (1,N)
 
     def transform(self, model, node):
-        dim = len(node.get_input_variable().shape) - 1        
+        dim = len(node.get_input_variable().shape) - 1
         input_shape = node.get_input_variable().shape
 
         pointwise_attrs = {
@@ -22,37 +26,41 @@ def transform(self, model, node):
         }
 
         if dim == 1:
-            pointwise_attrs.update({
-                'in_width': input_shape[0],
-                'out_width': input_shape[0],
-                'filt_width': 1,
-                'stride_width': 1,
-                'pad_left': 0,
-                'pad_right': 0,
-            })
+            pointwise_attrs.update(
+                {
+                    'in_width': input_shape[0],
+                    'out_width': input_shape[0],
+                    'filt_width': 1,
+                    'stride_width': 1,
+                    'pad_left': 0,
+                    'pad_right': 0,
+                }
+            )
         elif dim == 2:
-            pointwise_attrs.update({
-                'in_height': input_shape[0],
-                'in_width': input_shape[1],
-                'out_height': input_shape[0],
-                'out_width': input_shape[1],
-                'filt_height': 1,
-                'filt_width': 1,
-                'stride_height': 1,
-                'stride_width': 1,
-                'pad_top': 0,
-                'pad_bottom': 0,
-                'pad_left': 0,
-                'pad_right': 0,
-            })
+            pointwise_attrs.update(
+                {
+                    'in_height': input_shape[0],
+                    'in_width': input_shape[1],
+                    'out_height': input_shape[0],
+                    'out_width': input_shape[1],
+                    'filt_height': 1,
+                    'filt_width': 1,
+                    'stride_height': 1,
+                    'stride_width': 1,
+                    'pad_top': 0,
+                    'pad_bottom': 0,
+                    'pad_left': 0,
+                    'pad_right': 0,
+                }
+            )
         else:
             raise Exception('Cannot replace Dense over {dim}D tensor with Conv{dim}D.'.format(dim=dim))
 
         class_name = 'PointwiseConv' + str(dim) + 'D'
         pw_node = model.make_node(class_name, node.name, pointwise_attrs, node.inputs.copy())
-        if len(node.weights['weight'].data.shape) == 2: # This can happen if we assign weights of Dense layer to 1x1 Conv2D
-            pw_node.weights['weight'].data = np.expand_dims(node.weights['weight'].data, axis=(0,1))
+        if len(node.weights['weight'].data.shape) == 2:  # This can happen if we assign weights of Dense layer to 1x1 Conv2D
+            pw_node.weights['weight'].data = np.expand_dims(node.weights['weight'].data, axis=tuple(range(dim)))
         pw_node.weights['bias'].data = node.weights['bias'].data
         model.replace_node(node, pw_node)
-        
+
         return True
