Reshape fixes: don't repack stream for flatten; remove final reshape

hls4ml/model/optimizer/__init__.py

@@ -12,7 +12,7 @@
 from hls4ml.model.optimizer.passes.conv_same_pad import InsertZeroPaddingBeforeConv2D
 from hls4ml.model.optimizer.passes.pointwise import OptimizePointwiseConv
 from hls4ml.model.optimizer.passes.clone import CloneOutput
-from hls4ml.model.optimizer.passes.repack_stream import ReshapeStream, BroadcastStream
+from hls4ml.model.optimizer.passes.repack_stream import ReshapeStream, BroadcastStream, RemoveFinalReshape
 from hls4ml.model.optimizer.passes.transpose_opt import RemoveUselessTranspose
 from hls4ml.model.optimizer.passes.multi_dense import ReplaceMultidimensionalDenseWithConv
 
@@ -40,6 +40,7 @@
 register_pass('conv2d_same_pad', InsertZeroPaddingBeforeConv2D)
 register_pass('optimize_pointwise_conv', OptimizePointwiseConv)
 register_pass('clone_output', CloneOutput)
+register_pass('remove_final_reshape', RemoveFinalReshape)
 register_pass('reshape_stream', ReshapeStream)
 register_pass('remove_useless_transpose', RemoveUselessTranspose)
 register_pass('replace_multidense_conv', ReplaceMultidimensionalDenseWithConv)

hls4ml/model/optimizer/passes/repack_stream.py

@@ -71,7 +71,8 @@ def config_cpp(self):
 class ReshapeStream(OptimizerPass):
     ''' Repacks stream for Reshape layer '''
     def match(self, node):
-        return node.__class__.__name__ == 'Reshape'
+        # do not run optimizer pass for a flatten layer (1 output dimension)
+        return node.__class__.__name__ == 'Reshape' and len(node.get_output_variable().shape) > 1
 
     def transform(self, model, node):
         if model.config.backend.name not in ['Vivado', 'VivadoAccelerator'] or \
@@ -121,3 +122,19 @@ def transform(self, model, node):
         node.inputs[idx] = brdcst_out
 
         return True
+
+class RemoveFinalReshape(OptimizerPass):
+    ''' Remove reshape if final layer '''
+    def match(self, node):
+        # match if reshape is final node
+        return node.__class__.__name__ == 'Reshape' and not node.get_output_nodes()
+
+    def transform(self, model, node):
+        if model.config.get_config_value('IOType') == 'io_parallel':
+            print('WARNING: Final layer is a Reshape, which does not affect the output for io_parallel; removing it')
+            # remove, but don't rewire because it's the output layer
+            model.remove_node(node, rewire=False) 
+            return True
+        elif model.config.get_config_value('IOType') == 'io_stream':
+            print('WARNING: Final layer is a Reshape, which may incur a large resource cost for io_stream; consider removing it')
+        return False

test/pytest/test_graph.py

@@ -1,6 +1,7 @@
 import hls4ml
 import numpy as np
 import pytest
+import tensorflow as tf
 
 class Reader:
     def get_weights_data(self, name, var):
@@ -107,3 +108,34 @@ def test_graph_branch(iotype, batch):
   y = model.predict([X0, X1]).reshape(y_expected.shape)
   # check the output
   np.testing.assert_allclose(y, y_expected, rtol=1, atol=2**-16)
+
+@pytest.mark.parametrize('iotype', ['io_parallel', 'io_stream'])
+def test_final_reshape(iotype):
+  ''' Test case for a model with a Reshape as the final layer '''
+  inputs = tf.keras.layers.Input(shape=(1,1,1)) # 1 input pixel
+  conv = tf.keras.layers.Conv2D(6,1) # 6 filters, 1x1 kernel
+  x = conv(inputs)
+  conv.set_weights([np.linspace(1,6,6).reshape(1,1,1,6), np.zeros(6)]) # ascending int weights, 0 bias
+  x = tf.keras.layers.Reshape((3,2))(x) # reshape the (1,1,6) output to (3,2)
+  model = tf.keras.models.Model(inputs=inputs, outputs=x)
+
+  # create the HLSModel
+  config = hls4ml.utils.config_from_keras_model(model, granularity='model')
+  hls_model = hls4ml.converters.convert_from_keras_model(model,
+                                                         output_dir=f'hls4mlprj_graph_final_reshape_{iotype}',
+                                                         backend='Vivado',
+                                                         io_type = iotype,
+                                                         hls_config=config)
+  hls_model.compile()
+
+  # Test on ascending integers. The weights mean that each output pixel/neuron has
+  # a different value
+  X = np.linspace(-4,4,9).reshape(9,1,1,1)
+  y = model.predict(X)
+  y_hls = hls_model.predict(X).reshape(y.shape)
+  # because of integer inputs and integer weights, we can expect exact matching
+  np.testing.assert_allclose(y, y_hls, rtol=0)
+
+
+
+