Merge pull request #794 from vloncar/nested_model

Support for parsing nested models

Author: jmitrevs

hls4ml/converters/keras/model.py

@@ -0,0 +1,45 @@
+from hls4ml.converters.keras_to_hls import (
+    KerasFileReader,
+    KerasModelReader,
+    KerasNestedFileReader,
+    keras_handler,
+    parse_default_keras_layer,
+    parse_keras_model,
+)
+
+model_layers = ['Sequential', 'Functional']
+
+
+@keras_handler(*model_layers)
+def parse_model_layer(keras_layer, input_names, input_shapes, data_reader):
+    assert keras_layer['class_name'] in model_layers
+
+    layer = parse_default_keras_layer(keras_layer, input_names)
+    layer['class_name'] = 'LayerGroup'
+
+    if isinstance(data_reader, KerasNestedFileReader):
+        # In the .h5 file, the paths don't go more than one level deep
+        nested_path = data_reader.nested_path
+    else:
+        nested_path = layer['name']
+
+    if isinstance(data_reader, KerasFileReader):
+        nested_reader = KerasNestedFileReader(data_reader, nested_path)
+    else:
+        nested_reader = KerasModelReader(data_reader.model.get_layer(layer['name']))
+
+    layer_list, input_layers, output_layers, output_shapes = parse_keras_model(keras_layer, nested_reader)
+
+    if output_layers is None:
+        last_layer = layer_list[-1]['name']
+    else:
+        last_layer = output_layers[0]
+    output_shape = output_shapes[last_layer]
+
+    layer['layer_list'] = layer_list
+    layer['input_layers'] = input_layers if input_layers is not None else []
+    layer['output_layers'] = output_layers if output_layers is not None else []
+    layer['data_reader'] = nested_reader
+    layer['output_shape'] = output_shape
+
+    return layer, output_shape

hls4ml/converters/keras_to_hls.py

@@ -47,8 +47,7 @@ def get_weights_data(self, layer_name, var_name):
 
 class KerasNestedFileReader(KerasFileReader):
     def __init__(self, data_reader, nested_path):
-        self.config = data_reader.config
-        self.h5file = h5py.File(self.config['KerasH5'], mode='r')
+        super().__init__(data_reader.config)
         self.nested_path = nested_path
 
     def _find_data(self, layer_name, var_name):
@@ -319,18 +318,19 @@ def parse_keras_model(model_arch, reader):
             inputs_map[layer['name']] = act_layer['name']
             if output_layers is not None and layer['name'] in output_layers:
                 output_layers = [act_layer['name'] if name == layer['name'] else name for name in output_layers]
+            output_shapes[act_layer['name']] = output_shape
             layer_list.append(act_layer)
 
         assert output_shape is not None
 
         output_shapes[layer['name']] = output_shape
 
-    return layer_list, input_layers, output_layers
+    return layer_list, input_layers, output_layers, output_shapes
 
 
 def keras_to_hls(config):
     model_arch, reader = get_model_arch(config)
-    layer_list, input_layers, output_layers = parse_keras_model(model_arch, reader)
+    layer_list, input_layers, output_layers, _ = parse_keras_model(model_arch, reader)
     print('Creating HLS model')
     hls_model = ModelGraph(config, layer_list, input_layers, output_layers)
     return hls_model

hls4ml/model/layers.py

@@ -1273,6 +1273,24 @@ def _initialize_transforms(self):
         self._output_features = self.attributes['n_out_features'][-1]
 
 
+class LayerGroup(Layer):
+    _expected_attributes = [
+        Attribute('layer_list', value_type=list),
+        Attribute('input_layers', value_type=list),
+        Attribute('output_layers', value_type=list),
+        Attribute('data_reader', value_type=object),
+        Attribute('output_shape', value_type=list),
+    ]
+
+    def initialize(self):
+        shape = self.get_attr('output_shape')
+        if shape[0] is None:
+            shape.pop(0)
+        dims = [f'N_INPUT_{self.index}_{i+1}' for i in range(len(shape))]
+
+        self.add_output_variable(shape, dims)
+
+
 layer_map = {
     'Input': Input,
     'InputLayer': Input,
@@ -1324,6 +1342,7 @@ def _initialize_transforms(self):
     'GRU': GRU,
     'GarNet': GarNet,
     'GarNetStack': GarNetStack,
+    'LayerGroup': LayerGroup,
     # TensorFlow-specific layers:
     'BiasAdd': BiasAdd,
 }

hls4ml/model/optimizer/__init__.py

@@ -36,6 +36,7 @@
         'channels_last_converter',
         'fuse_bias_add',
         'remove_useless_transpose',
+        'expand_layer_group',
         'output_rounding_saturation_mode',
         'qkeras_factorize_alpha',
         'extract_ternary_threshold',

hls4ml/model/optimizer/passes/expand_layer_group.py

@@ -0,0 +1,46 @@
+from hls4ml.model.layers import Input, LayerGroup
+from hls4ml.model.optimizer import OptimizerPass
+
+
+class ExpandLayerGroup(OptimizerPass):
+    '''Expands LayerGroup (a nested model) into the parent model.'''
+
+    def match(self, node):
+        return isinstance(node, LayerGroup)
+
+    def transform(self, model, node):
+        layer_list = node.get_attr('layer_list')
+
+        # We'll keep track of inserted Input nodes to remove later
+        inserted_input_nodes = []
+
+        for i, layer in enumerate(layer_list):
+            kind = layer['class_name']
+            name = layer['name']
+            inputs = layer.get('inputs', [])
+            outputs = layer.get('outputs', [])
+
+            if name in model.graph.keys():
+                raise Exception(f'Layer names must be unique: "{name}" already found in the model graph.')
+
+            if len(inputs) == 0:
+                if kind in ['InputLayer', 'Input']:
+                    inputs = node.inputs.copy()
+                else:
+                    inputs = model.graph[layer_list[i - 1]['name']].outputs.copy()
+            if len(outputs) == 0:
+                outputs = [name]
+
+            new_node = model.make_node(kind, name, layer, inputs, outputs)
+            model.insert_node(new_node)
+            if isinstance(new_node, Input):
+                inserted_input_nodes.append(new_node)
+
+        rewire = not node.outputs[0] in model.outputs
+
+        model.remove_node(node, rewire)
+
+        for input_node in inserted_input_nodes:
+            model.remove_node(input_node, rewire=True)
+
+        return True

hls4ml/utils/config.py

@@ -133,7 +133,7 @@ def config_from_keras_model(
 
     reader = hls4ml.converters.KerasModelReader(model)
 
-    layer_list, _, _ = hls4ml.converters.parse_keras_model(model_arch, reader)
+    layer_list, _, _, _ = hls4ml.converters.parse_keras_model(model_arch, reader)
 
     def make_layer_config(layer):
         cls_name = layer['class_name']

test/pytest/test_keras_nested_model.py

@@ -0,0 +1,173 @@
+""" Test that nested models in Keras is properly parsed and expanded by the optimizers.
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+from tensorflow.keras.layers import Dense, Input
+from tensorflow.keras.models import Model, Sequential
+
+import hls4ml
+
+test_root_path = Path(__file__).parent
+
+
+def make_nested_model(input_shape):
+    """
+    This model will have the following architecture:
+    Functional (fun_model)
+        Dense (fun_first_dense)
+        Sequential (seq_sub)
+            Dense
+            Dense
+        Dense (fun_middle_dense)
+        Functional (fun_sub)
+            Dense
+            Dense
+        Dense (fun_last_dense)
+    """
+    seq_sub = Sequential(name='seq_sub')
+    seq_sub.add(Dense(5, activation='linear', input_shape=(5,), name='seq_sub_dense_1'))
+    seq_sub.add(Dense(3, activation='linear', name='seq_sub_dense_2'))
+
+    fun_input = Input(shape=(8,), name='fun_input')
+    fun_x = Dense(7, activation='linear', name='fun_sub_dense_1')(fun_input)
+    fun_x = Dense(6, activation='linear', name='fun_sub_dense_2')(fun_x)
+    fun_sub = Model(inputs=fun_input, outputs=fun_x, name='fun_sub')
+
+    input = Input(shape=input_shape, name='model_input')
+    x = Dense(5, activation='linear', name='fun_first_dense')(input)
+    x = seq_sub(x)
+    x = Dense(8, activation='linear', name='fun_middle_dense')(x)
+    x = fun_sub(x)
+    x = Dense(4, activation='linear', name='fun_last_dense')(x)
+    fun_model = Model(inputs=input, outputs=x, name='fun_model')
+
+    return fun_model
+
+
+def make_sub_nested_model(input_shape):
+    """
+    The following abomination will create this hierarchy:
+    Sequential
+        Dense (first_dense)
+        Functional (fun_model)
+            Dense (fun_first_dense)
+            Sequential (fun_model_seq_sub)
+                Dense
+                Dense
+            Dense (fun_middle_dense)
+            Functional (fun_model_fun_sub)
+                Dense
+                Dense
+            Dense (fun_last_dense)
+        Dense (middle_dense)
+        Sequential (seq_model)
+            Dense
+            Functional (seq_model_fun_sub)
+                Dense
+                Dense
+            Dense
+            Sequential (seq_model_seq_sub)
+                Dense
+                Dense
+            Dense
+        Dense (last_dense)
+    """
+    fun_model_seq_sub = Sequential(name='fun_model_seq_sub')
+    fun_model_seq_sub.add(Dense(5, activation='linear', input_shape=(5,), name='fun_seq_sub_dense_1'))
+    fun_model_seq_sub.add(Dense(3, activation='linear', name='fun_seq_sub_dense_2'))
+
+    fun_fun_input = Input(shape=(8,), name='fun_fun_input')
+    fun_fun_x = Dense(7, activation='linear', name='fun_fun_sub_dense_1')(fun_fun_input)
+    fun_fun_x = Dense(6, activation='linear', name='fun_fun_sub_dense_2')(fun_fun_x)
+    fun_model_fun_sub = Model(inputs=fun_fun_input, outputs=fun_fun_x, name='fun_model_fun_sub')
+
+    fun_input = Input(shape=(10,), name='fun_input')
+    fun_x = Dense(5, activation='linear', name='fun_first_dense')(fun_input)
+    fun_x = fun_model_seq_sub(fun_x)
+    fun_x = Dense(8, activation='linear', name='fun_middle_dense')(fun_x)
+    fun_x = fun_model_fun_sub(fun_x)
+    fun_x = Dense(4, activation='linear', name='fun_last_dense')(fun_x)
+    fun_model = Model(inputs=fun_input, outputs=fun_x, name='fun_model')
+
+    seq_fun_input = Input(shape=(2,), name='seq_fun_input')
+    seq_fun_x = Dense(9, activation='linear', name='seq_fun_sub_dense_1')(seq_fun_input)
+    seq_fun_x = Dense(3, activation='linear', name='seq_fun_sub_dense_2')(seq_fun_x)
+    seq_model_fun_sub = Model(inputs=seq_fun_input, outputs=seq_fun_x, name='seq_model_fun_sub')
+
+    seq_model_seq_sub = Sequential(name='seq_model_seq_sub')
+    seq_model_seq_sub.add(Dense(5, activation='linear', input_shape=(2,), name='seq_seq_sub_dense_1'))
+    seq_model_seq_sub.add(Dense(7, activation='linear', name='seq_seq_sub_dense_2'))
+
+    seq_model = Sequential(name='seq_model')
+    seq_model.add(Dense(2, activation='linear', input_shape=(6,), name='seq_first_dense'))
+    seq_model.add(seq_model_fun_sub)
+    seq_model.add(Dense(2, activation='linear', name='seq_middle_dense'))
+    seq_model.add(seq_model_seq_sub)
+    seq_model.add(Dense(2, activation='linear', name='seq_last_dense'))
+
+    model = Sequential()
+    model.add(Dense(10, activation='linear', input_shape=input_shape, name='first_dense'))
+    model.add(fun_model)
+    model.add(Dense(6, activation='linear', name='middle_dense'))
+    model.add(seq_model)
+    model.add(Dense(4, activation='linear', name='last_dense'))
+
+    return model
+
+
+def randX(batch_size, N):
+    return np.random.rand(batch_size, N)
+
+
+@pytest.fixture(scope='module')
+def randX_20_15():
+    return randX(20, 15)
+
+
+@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
+def test_nested_model(randX_20_15, backend, io_type):
+    n_in = 15
+    input_shape = (n_in,)
+    keras_model = make_nested_model(input_shape)
+    keras_model.compile(optimizer='adam', loss='mae')
+
+    config = hls4ml.utils.config_from_keras_model(keras_model, default_precision='fixed<24,12>')
+    prj_name = f'hls4mlprj_nested_model_{backend}_{io_type}'
+    output_dir = str(test_root_path / prj_name)
+    hls_model = hls4ml.converters.convert_from_keras_model(
+        keras_model, hls_config=config, output_dir=output_dir, io_type=io_type, backend=backend
+    )
+    hls_model.compile()
+
+    X = randX_20_15
+    y_keras = keras_model.predict(X)
+    y_hls4ml = hls_model.predict(X)
+
+    np.testing.assert_allclose(y_keras.ravel(), y_hls4ml.ravel(), rtol=1e-2, atol=0.02)
+
+
+@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
+def test_sub_nested_model(randX_20_15, backend, io_type):
+    n_in = 15
+    input_shape = (n_in,)
+    keras_model = make_sub_nested_model(input_shape)
+    keras_model.compile(optimizer='adam', loss='mae')
+
+    config = hls4ml.utils.config_from_keras_model(keras_model, default_precision='fixed<24,12>')
+    prj_name = f'hls4mlprj_sub_nested_model_{backend}_{io_type}'
+    output_dir = str(test_root_path / prj_name)
+    hls_model = hls4ml.converters.convert_from_keras_model(
+        keras_model, hls_config=config, output_dir=output_dir, io_type=io_type, backend=backend
+    )
+    hls_model.compile()
+
+    X = randX_20_15
+    y_keras = keras_model.predict(X)
+    y_hls4ml = hls_model.predict(X)
+
+    np.testing.assert_allclose(y_keras.ravel(), y_hls4ml.ravel(), rtol=1e-2, atol=0.02)