ADD Bidirectional layers support

Author: Enrico Lupi

hls4ml/backends/fpga/fpga_backend.py

@@ -14,8 +14,7 @@
     Activation,
     BatchNormalization,
     BatchNormOnnx,
-    BidirectionalGRU,
-    BidirectionalLSTM,
+    Bidirectional,
     Conv,
     Conv1D,
     Conv2D,
@@ -70,8 +69,7 @@ def __init__(self, name):
             SimpleRNN,
             LSTM,
             GRU,
-            BidirectionalLSTM,
-            BidirectionalGRU,
+            Bidirectional,
             Dot,
             Conv,
             MatMul,
@@ -217,34 +215,30 @@ def get_layer_mult_size(self, layer):
             n_out = layer.get_attr('n_filt')
             return n_in, n_out
 
-        if 'BidirectionalLSTM' in layer.class_name:
-            n_in = layer.get_attr('n_in')
-            n_out = layer.get_attr('n_out') * 2  # /2*4
-            n_in_recr = layer.get_attr('n_out') // 2
-            n_out_recr = n_out
-            return n_in, n_out, n_in_recr, n_out_recr
-
         if 'LSTM' in layer.class_name:
             n_in = layer.get_attr('n_in')
             n_out = layer.get_attr('n_out') * 4
             n_in_recr = layer.get_attr('n_out')
             n_out_recr = n_out
             return n_in, n_out, n_in_recr, n_out_recr
 
-        if 'BidirectionalGRU' in layer.class_name:
-            n_in = layer.get_attr('n_in')
-            n_out = layer.get_attr('n_out') // 2 * 3
-            n_in_recr = layer.get_attr('n_out') // 2
-            n_out_recr = n_out
-            return n_in, n_out, n_in_recr, n_out_recr
-
         if 'GRU' in layer.class_name:
             n_in = layer.get_attr('n_in')
             n_out = layer.get_attr('n_out') * 3
             n_in_recr = layer.get_attr('n_out')
             n_out_recr = n_out
             return n_in, n_out, n_in_recr, n_out_recr
 
+        if 'Bidirectional' in layer.class_name:
+            result = []
+            for d in ['forward', 'backward']:
+                n_in = layer.get_attr('n_in')
+                n_out = layer.get_attr(f'{d}_n_states') * 3
+                n_in_recr = layer.get_attr(f'{d}_n_states')
+                n_out_recr = n_out
+                result.append((n_in, n_out, n_in_recr, n_out_recr))
+            return result
+
         raise Exception(f'Cannot get mult size for layer {layer.name} ({layer.class_name})')
 
     def get_valid_reuse_factors(self, n_in, n_out):
@@ -295,6 +289,7 @@ def set_closest_reuse_factor(self, layer, n_in, n_out, attribute='reuse_factor',
         if not include_max_rf:
             valid_rf.pop()
         chosen_rf = layer.get_attr(attribute)
+        print("\n\nREuse factor:", chosen_rf, "\n\n")
         if chosen_rf not in valid_rf:
             closest_rf = self.get_closest_reuse_factor(valid_rf, chosen_rf)
             valid_rf_str = ','.join(map(str, valid_rf))

hls4ml/backends/vivado/passes/recurrent_templates.py

@@ -3,8 +3,7 @@
 from hls4ml.model.layers import (
     GRU,
     LSTM,
-    BidirectionalGRU,
-    BidirectionalLSTM,
+    Bidirectional,
     Conv1D,
     Conv2D,
     Dense,
@@ -18,9 +17,7 @@ class ApplyResourceStrategy(OptimizerPass):
     '''Transposes the weights to use the dense_resource matrix multiply routine'''
 
     def match(self, node):
-        node_matches = isinstance(
-            node, (Dense, Conv1D, SeparableConv1D, Conv2D, SeparableConv2D, LSTM, GRU, BidirectionalLSTM, BidirectionalGRU)
-        )
+        node_matches = isinstance(node, (Dense, Conv1D, SeparableConv1D, Conv2D, SeparableConv2D, LSTM, GRU, Bidirectional))
         is_resource_strategy = node.get_attr('strategy', '').lower() in ['resource', 'resource_unrolled']
         already_transformed = node.get_attr('_weights_transposed', False) is True
         return node_matches and is_resource_strategy and not already_transformed
@@ -48,11 +45,10 @@ def transform(self, model, node):
             node.weights['pointwise'].data = np.transpose(
                 node.weights['pointwise'].data, axes=[3, 0, 1, 2]
             )  # (H,W,C,F) => (F,H,W,C)
-        elif isinstance(node, (BidirectionalLSTM, BidirectionalGRU)):
-            node.weights['weight'].data = np.transpose(node.weights['weight'].data)
-            node.weights['recurrent_weight'].data = np.transpose(node.weights['recurrent_weight'].data)
-            node.weights['weight_b'].data = np.transpose(node.weights['weight_b'].data)
-            node.weights['recurrent_weight_b'].data = np.transpose(node.weights['recurrent_weight_b'].data)
+        elif isinstance(node, (Bidirectional)):
+            for d in ['forward', 'backward']:
+                node.weights[f'{d}_weight'].data = np.transpose(node.weights[f'{d}_weight'].data)
+                node.weights[f'{d}_recurrent_weight'].data = np.transpose(node.weights[f'{d}_recurrent_weight'].data)
         elif isinstance(node, (LSTM, GRU)):
             node.weights['weight'].data = np.transpose(node.weights['weight'].data)
             node.weights['recurrent_weight'].data = np.transpose(node.weights['recurrent_weight'].data)

hls4ml/backends/vivado/passes/resource_strategy.py

@@ -12,8 +12,7 @@
 from hls4ml.model.layers import (
     GRU,
     LSTM,
-    BidirectionalGRU,
-    BidirectionalLSTM,
+    Bidirectional,
     Conv1D,
     Conv2D,
     Dense,
@@ -48,7 +47,7 @@ def __init__(self):
 
     def _register_layer_attributes(self):
         # Add RNN-specific attributes, recurrent_reuse_factor and static implementation
-        rnn_layers = [SimpleRNN, LSTM, GRU, BidirectionalLSTM, BidirectionalGRU]
+        rnn_layers = [SimpleRNN, LSTM, GRU]
 
         for layer in rnn_layers:
             attrs = self.attribute_map.get(layer, [])
@@ -60,6 +59,24 @@ def _register_layer_attributes(self):
             attrs.append(TypeAttribute('table', default=FixedPrecisionType(18, 8), description=descriptions.table_type))
             self.attribute_map[layer] = attrs
 
+        bidir_rnn_layers = [Bidirectional]
+        for layer in bidir_rnn_layers:
+            attrs = self.attribute_map.get(layer, [])
+            attrs.append(ConfigurableAttribute('forward_reuse_factor', default=1, description=descriptions.reuse_factor))
+            attrs.append(ConfigurableAttribute('backward_reuse_factor', default=1, description=descriptions.reuse_factor))
+            attrs.append(
+                ConfigurableAttribute('forward_recurrent_reuse_factor', default=1, description=descriptions.reuse_factor)
+            )
+            attrs.append(
+                ConfigurableAttribute('backward_recurrent_reuse_factor', default=1, description=descriptions.reuse_factor)
+            )
+            attrs.append(
+                ConfigurableAttribute('static', value_type=bool, default=True, description=descriptions.recurrent_static)
+            )
+            attrs.append(ConfigurableAttribute('table_size', default=1024, description=descriptions.table_size))
+            attrs.append(TypeAttribute('table', default=FixedPrecisionType(18, 8), description=descriptions.table_type))
+            self.attribute_map[layer] = attrs
+
         # Add ParallelizationFactor to Conv1D/2D
         pf_layers = [
             Conv1D,
@@ -657,6 +674,45 @@ def init_time_distributed(self, layer):
             warn(f'Cannot unroll time step loop in layer "{layer.name}" while using "io_stream".')
             loop_mode = 'off'
         layer.set_attr('time_step_loop_parallelism', loop_mode)
+        
+    @layer_optimizer(Bidirectional)
+    def init_bidirectional(self, layer):
+        reuse_factor = layer.model.config.get_reuse_factor(layer)
+
+        for i, d in enumerate(['forward', 'backward']):
+            layer.set_attr(f'{d}_reuse_factor', reuse_factor)
+            layer.set_attr(f'{d}_recurrent_reuse_factor', reuse_factor)
+
+            if layer.model.config.is_resource_strategy(layer):
+                n_in, n_out, n_in_recr, n_out_recr = self.get_layer_mult_size(layer)[i]
+                self.set_closest_reuse_factor(layer, n_in, n_out, attribute=f'{d}_reuse_factor')
+                self.set_closest_reuse_factor(layer, n_in_recr, n_out_recr, attribute=f'{d}_recurrent_reuse_factor')
+                layer.set_attr('strategy', 'resource')
+
+            elif layer.model.config.get_strategy(layer).lower() == 'resource_unrolled':
+                use_resource_instead = False
+                if layer.get_attr('reuse_factor', 1) == 1:
+                    print(
+                        f'Unrolled resource strategy cannot be combined with reuse factor 1 in layer "{layer.name} ({d})". '
+                        'Using "resource" strategy instead.'
+                    )
+                use_resource_instead = True
+
+                n_in, n_out, n_in_recr, n_out_recr = self.get_layer_mult_size(layer)[i]
+                if use_resource_instead:
+                    self.set_closest_reuse_factor(layer, n_in, n_out, attribute=f'{d}_reuse_factor')
+                    self.set_closest_reuse_factor(layer, n_in_recr, n_out_recr, attribute=f'{d}_recurrent_reuse_factor')
+                    layer.set_attr('strategy', 'resource')
+                else:
+                    self.set_closest_reuse_factor(layer, n_in, n_out, attribute=f'{d}_reuse_factor', include_max_rf=False)
+                    self.set_closest_reuse_factor(
+                        layer, n_in_recr, n_out_recr, attribute=f'{d}_recurrent_reuse_factor', include_max_rf=False
+                    )
+                    layer.set_attr('strategy', 'resource_unrolled')
+            else:
+                layer.set_attr('strategy', 'latency')
+
+        layer.set_attr('index_t', NamedType(f'layer{layer.index}_index', IntegerPrecisionType(width=1, signed=False)))
 
     @layer_optimizer(GarNet)
     def init_garnet(self, layer):

hls4ml/backends/vivado/vivado_backend.py

@@ -241,7 +241,7 @@ def parse_keras_model(model_arch, reader):
         'HGQ>UnaryLUT',
     ]
     # Recurrent layers
-    recurrent_layers = ['SimpleRNN', 'LSTM', 'GRU', 'QSimpleRNN', 'QLSTM', 'QGRU', 'BidirectionalLSTM', 'BidirectionalGRU']
+    recurrent_layers = ['SimpleRNN', 'LSTM', 'GRU', 'QSimpleRNN', 'QLSTM', 'QGRU', 'Bidirectional']
     # All supported layers
     supported_layers = get_supported_keras_layers() + skip_layers
 

hls4ml/converters/keras_v2_to_hls.py

@@ -1402,40 +1402,6 @@ def initialize(self):
             self.add_weights_variable(name='recurrent_bias', var_name='br{index}', data=recurrent_bias)
 
 
-class BidirectionalLSTM(LSTM):
-    _expected_attributes = [
-        WeightAttribute('weight_b'),
-        WeightAttribute('bias_b'),
-        WeightAttribute('recurrent_weight_b'),
-        WeightAttribute('recurrent_bias_b'),
-        TypeAttribute('weight_b'),
-        TypeAttribute('bias_b'),
-        TypeAttribute('recurrent_weight_b'),
-        TypeAttribute('recurrent_bias_b'),
-        ChoiceAttribute('merge_mode', ['sum', 'mul', 'concat', 'ave'], configurable=False, default='concat'),
-    ]
-
-    def initialize(self):
-        super().initialize()
-
-        # Add backward layer parameters
-        # weights
-        self.add_weights_variable(name='weight_b', var_name='w_b{index}')
-
-        # recurrent weights
-        self.add_weights_variable(name='recurrent_weight_b', var_name='wr_b{index}')
-
-        # biases
-        self.add_weights_variable(name='bias_b', var_name='b_b{index}')
-
-        if "pytorch" in self.attributes.keys():
-            self.add_weights_variable(name='recurrent_bias_b', var_name='br_b{index}')
-        else:
-            recurrent_weight_b = self.get_attr('recurrent_weight_b_data')
-            recurrent_bias_b = np.zeros(recurrent_weight_b.shape[1])
-            self.add_weights_variable(name='recurrent_bias_b', var_name='br_b{index}', data=recurrent_bias_b)
-
-
 class GRU(Layer):
     _expected_attributes = [
         Attribute('n_out'),
@@ -1509,34 +1475,6 @@ def initialize(self):
 
         self.add_output_variable(shape, dims)
 
-        
-class BidirectionalGRU(GRU):
-    _expected_attributes = [
-        WeightAttribute('weight_b'),
-        WeightAttribute('bias_b'),
-        WeightAttribute('recurrent_weight_b'),
-        WeightAttribute('recurrent_bias_b'),
-        TypeAttribute('weight_b'),
-        TypeAttribute('bias_b'),
-        TypeAttribute('recurrent_weight_b'),
-        TypeAttribute('recurrent_bias_b'),
-        ChoiceAttribute('merge_mode', ['sum', 'mul', 'concat', 'ave'], configurable=False, default='concat'),
-    ]
-
-    def initialize(self):
-        super().initialize()
-
-        # Add backward layer parameters
-        # weights
-        self.add_weights_variable(name='weight_b', var_name='w_b{index}')
-
-        # recurrent weights
-        self.add_weights_variable(name='recurrent_weight_b', var_name='wr_b{index}')
-
-        # biases
-        self.add_weights_variable(name='bias_b', var_name='b_b{index}')
-        self.add_weights_variable(name='recurrent_bias_b', var_name='br_b{index}')
-
 
 class Bidirectional(Layer):
     _expected_attributes = [
@@ -1609,9 +1547,7 @@ def initialize(self):
                         name=f'{dir}_recurrent_bias', var_name=(f'br_{dir[0]}_' + '{index}'), data=recurrent_bias
                     )
             else:
-                self.add_weights_variable(
-                    name=f'{dir}_recurrent_bias', var_name=(f'br_{dir[0]}_' + '{index}'), data=recurrent_bias
-                )
+                self.add_weights_variable(name=f'{dir}_recurrent_bias', var_name=(f'br_{dir[0]}_' + '{index}'))
 
 
 class GarNet(Layer):
@@ -1902,8 +1838,6 @@ def initialize(self):
     'SimpleRNN': SimpleRNN,
     'LSTM': LSTM,
     'GRU': GRU,
-    'BidirectionalLSTM': BidirectionalLSTM,
-    'BidirectionalGRU': BidirectionalGRU,
     'Bidirectional': Bidirectional,
     'QSimpleRNN': SimpleRNN,
     'QLSTM': LSTM,

hls4ml/model/layers.py

@@ -81,7 +81,7 @@ def _infer_precision(self, node, types_to_infer):
         if node_class in ['Embedding']:
             return self._infer_embedding_precision(node, types_to_infer)
 
-        if node_class in ['SimpleRNN', 'LSTM', 'GRU', 'BidirectionalLSTM', 'BidirectionalGRU']:
+        if node_class in ['SimpleRNN', 'LSTM', 'GRU', 'Bidirectional']:
             return self._infer_rnn_precision(node, types_to_infer)
 
         if node_class in ['ParametrizedActivation']:
@@ -554,8 +554,9 @@ def _infer_rnn_precision(self, node, types_to_infer):
 
         # for now just do the weights and leave the rest for the default catch
         rnn_weights = ('weight', 'bias', 'recurrent_weight', 'recurrent_bias')
-        if node.attributes['direction'] == 'bidirectional':
-            rnn_weights += ('weight_b', 'bias_b', 'recurrent_weight_b', 'recurrent_bias_b')
+        if node.class_name == 'Bidirectional':
+            rnn_weights = [direction + '_' + weight for direction in ['forward', 'backward'] for weight in rnn_weights]
+
         for weightvar in rnn_weights:
             if f'{weightvar}_t' in types_to_infer:
                 self._infer_default_type(node, f'{weightvar}_t')