Skip to content

Commit

Permalink
Reshape fixes: don't repack stream for flatten; remove final reshape (f…
Browse files Browse the repository at this point in the history
…astmachinelearning#443)

* fix 2 reshape issues: don't reshape streams for flatten and remove final reshape

* Add a test for a model with Reshape as the final layer

* swap

* only remove for io_parallel; warn for both io_parallel and io_stream

Co-authored-by: Sioni Summers <[email protected]>
  • Loading branch information
jmduarte and thesps authored Nov 9, 2021
1 parent add0cfc commit 0f26308
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 2 deletions.
3 changes: 2 additions & 1 deletion hls4ml/model/optimizer/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down
19 changes: 18 additions & 1 deletion hls4ml/model/optimizer/passes/repack_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 \
Expand Down Expand Up @@ -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
32 changes: 32 additions & 0 deletions test/pytest/test_graph.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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)




0 comments on commit 0f26308

Please sign in to comment.