diff --git a/hls4ml/model/graph.py b/hls4ml/model/graph.py index 670c12537f..43a924fbf9 100644 --- a/hls4ml/model/graph.py +++ b/hls4ml/model/graph.py @@ -591,6 +591,7 @@ def _get_top_function(self, x): xlist = [x] else: xlist = x + n_outputs = len(self.get_output_variables()) for xi in xlist: if not isinstance(xi, np.ndarray): @@ -610,7 +611,7 @@ def _get_top_function(self, x): top_function.restype = None - top_function.argtypes = [npc.ndpointer(ctype, flags="C_CONTIGUOUS") for i in range(len(xlist)+1)] + top_function.argtypes = [npc.ndpointer(ctype, flags="C_CONTIGUOUS") for i in range(len(xlist) + n_outputs)] return top_function, ctype @@ -637,6 +638,7 @@ def predict(self, x): top_function, ctype = self._get_top_function(x) n_samples = self._compute_n_samples(x) n_inputs = len(self.get_input_variables()) + n_outputs = len(self.get_output_variables()) curr_dir = os.getcwd() os.chdir(self.config.get_output_dir() + '/firmware') @@ -647,25 +649,29 @@ def predict(self, x): try: for i in range(n_samples): - predictions = np.zeros(self.get_output_variables()[0].size(), dtype=ctype) + predictions = [np.zeros(yj.size(), dtype=ctype) for yj in self.get_output_variables()] if n_inputs == 1: - top_function(x[i], predictions) + inp = [x[i]] else: inp = [xj[i] for xj in x] - argtuple = inp - argtuple += [predictions] - argtuple = tuple(argtuple) - top_function(*argtuple) + argtuple = inp + argtuple += predictions + argtuple = tuple(argtuple) + top_function(*argtuple) output.append(predictions) - #Convert to numpy array - output = np.asarray(output) + # Convert to list of numpy arrays (one for each output) + output = [np.asarray([output[i_sample][i_output] for i_sample in range(n_samples)]) for i_output in range(n_outputs)] finally: os.chdir(curr_dir) - - if n_samples == 1: + + if n_samples == 1 and n_outputs == 1: + return output[0][0] + elif n_outputs == 1: return output[0] + elif n_samples == 1: + return [output_i[0] for output_i in output] else: return output diff --git a/test/pytest/test_graph.py b/test/pytest/test_graph.py index a187e59290..a1f3d53403 100644 --- a/test/pytest/test_graph.py +++ b/test/pytest/test_graph.py @@ -175,3 +175,31 @@ def test_broadcast_stream(shapes, layer): y = model.predict([X1, X2]) y_hls = hls_model.predict([X1, X2]).reshape(y.shape) np.testing.assert_allclose(y, y_hls, rtol=0) + +@pytest.mark.parametrize('batch', [1, 32]) +def test_multiple_outputs(batch): + ''' Test case for multple outputs ''' + input1 = tf.keras.layers.Input(shape=(10,)) + inputs = [input1] + output1 = tf.keras.layers.Dense(5, kernel_initializer='ones', use_bias=False)(input1) + output2 = tf.keras.layers.Dense(2, kernel_initializer='ones', use_bias=False)(input1) + outputs = [output1, output2] + model = tf.keras.models.Model(inputs=inputs, outputs=outputs) + + # create the ModelGraph + config = hls4ml.utils.config_from_keras_model(model, granularity='model', default_precision='ap_fixed<32,16>') + odir = str(test_root_path / 'hls4mlprj_graph_multiple_outputs') + hls_model = hls4ml.converters.convert_from_keras_model(model, + output_dir=odir, + backend='Vivado', + io_type='io_serial', + hls_config=config) + hls_model.compile() + + # Test with integers (for exact agreement) + X1 = np.random.randint(0, 100, size=(batch, 10)).astype(float) + y = model.predict(X1) + y_hls = hls_model.predict(X1) + for y_i, y_hls_i in zip(y, y_hls): + y_hls_i = y_hls_i.reshape(y_i.shape) + np.testing.assert_allclose(y_i, y_hls_i, rtol=0) \ No newline at end of file