Support keepdims in GlobalPooling layers (#716)

* Support keepdims in GlobalPooling layers

* Update ci-template.yml

* parametrize fixture for keepdims

* pre-commit

---------

Co-authored-by: Javier Duarte <jduarte@ucsd.edu>

Author: vloncar

hls4ml/converters/keras/pooling.py

@@ -67,14 +67,27 @@ def parse_global_pooling_layer(keras_layer, input_names, input_shapes, data_read
     assert 'Pooling' in keras_layer['class_name']
 
     layer = parse_default_keras_layer(keras_layer, input_names)
+    layer['keepdims'] = keras_layer['config']['keepdims']
 
     if int(layer['class_name'][-2]) == 1:
         (layer['n_in'], layer['n_filt']) = parse_data_format(input_shapes[0], layer['data_format'])
 
-        output_shape = [input_shapes[0][0], layer['n_filt']]
+        if layer['keepdims']:
+            if layer['data_format'] == 'channels_last':
+                output_shape = [input_shapes[0][0], 1, layer['n_filt']]
+            elif layer['data_format'] == 'channels_first':
+                output_shape = [input_shapes[0][0], layer['n_filt'], 1]
+        else:
+            output_shape = [input_shapes[0][0], layer['n_filt']]
     elif int(layer['class_name'][-2]) == 2:
         (layer['in_height'], layer['in_width'], layer['n_filt']) = parse_data_format(input_shapes[0], layer['data_format'])
 
-        output_shape = [input_shapes[0][0], layer['n_filt']]
+        if layer['keepdims']:
+            if layer['data_format'] == 'channels_last':
+                output_shape = [input_shapes[0][0], 1, 1, layer['n_filt']]
+            elif layer['data_format'] == 'channels_first':
+                output_shape = [input_shapes[0][0], layer['n_filt'], 1, 1]
+        else:
+            output_shape = [input_shapes[0][0], layer['n_filt']]
 
     return layer, output_shape

test/pytest/ci-template.yml

@@ -1,7 +1,7 @@
 .pytest:
   stage: test
-  image: gitlab-registry.cern.ch/fastmachinelearning/hls4ml-testing:0.2.base
-  tags: 
+  image: gitlab-registry.cern.ch/fastmachinelearning/hls4ml-testing:0.4.base
+  tags:
     - docker
   before_script:
     - source ~/.bashrc
@@ -14,10 +14,10 @@
   artifacts:
     when: always
     reports:
-      junit: 
+      junit:
         - test/pytest/report.xml
       coverage_report:
         coverage_format: cobertura
         path: test/pytest/coverage.xml
     paths:
-      - test/pytest/hls4mlprj*.tar.gz
+      - test/pytest/hls4mlprj*.tar.gz

test/pytest/test_globalpooling.py

@@ -1,94 +1,126 @@
+from pathlib import Path
+
+import numpy as np
 import pytest
+from tensorflow.keras.layers import GlobalAveragePooling1D, GlobalAveragePooling2D, GlobalMaxPooling1D, GlobalMaxPooling2D
 from tensorflow.keras.models import Sequential
-from tensorflow.keras.layers import GlobalAveragePooling1D, GlobalMaxPooling1D, GlobalAveragePooling2D, GlobalMaxPooling2D
-import numpy as np
+
 import hls4ml
-from pathlib import Path
 
 test_root_path = Path(__file__).parent
 
 in_shape = 18
 in_filt = 6
 atol = 5e-3
 
+
 @pytest.fixture(scope='module')
 def data_1d():
     return np.random.rand(100, in_shape, in_filt)
 
-@pytest.fixture(scope='module')
-def keras_model_max_1d():
-    model = Sequential()
-    model.add(GlobalMaxPooling1D(input_shape=(in_shape, in_filt)))
-    model.compile()
-    return model
 
 @pytest.fixture(scope='module')
-def keras_model_avg_1d():
+def keras_model_1d(request):
+    model_type = request.param['model_type']
+    keepdims = request.param['keepdims']
     model = Sequential()
-    model.add(GlobalAveragePooling1D(input_shape=(in_shape, in_filt)))
+    if model_type == 'avg':
+        model.add(GlobalAveragePooling1D(input_shape=(in_shape, in_filt), keepdims=keepdims))
+    elif model_type == 'max':
+        model.add(GlobalMaxPooling1D(input_shape=(in_shape, in_filt), keepdims=keepdims))
     model.compile()
-    return model
- 
+    return model, model_type, keepdims
+
 
 @pytest.mark.parametrize('backend', ['Quartus', 'Vivado'])
-@pytest.mark.parametrize('model_type', ['max', 'avg'])
+@pytest.mark.parametrize(
+    'keras_model_1d',
+    [
+        {'model_type': 'max', 'keepdims': True},
+        {'model_type': 'max', 'keepdims': False},
+        {'model_type': 'avg', 'keepdims': True},
+        {'model_type': 'avg', 'keepdims': False},
+    ],
+    ids=[
+        'model_type-max-keepdims-True',
+        'model_type-max-keepdims-False',
+        'model_type-avg-keepdims-True',
+        'model_type-avg-keepdims-False',
+    ],
+    indirect=True,
+)
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
-def test_global_pool1d(backend, keras_model_max_1d, keras_model_avg_1d, data_1d, model_type, io_type):
-    if model_type == 'avg':
-        model = keras_model_avg_1d
-    else:
-        model = keras_model_max_1d
-    
+def test_global_pool1d(backend, keras_model_1d, data_1d, io_type):
+
+    model, model_type, keepdims = keras_model_1d
+
     config = hls4ml.utils.config_from_keras_model(model, default_precision='ap_fixed<32,9>', granularity='name')
 
-    hls_model = hls4ml.converters.convert_from_keras_model(model,
-                                                           hls_config=config,
-                                                           io_type=io_type,
-                                                           output_dir=str(test_root_path / f'hls4mlprj_globalplool1d_{backend}_{io_type}_{model_type}'),
-                                                           backend=backend)
+    hls_model = hls4ml.converters.convert_from_keras_model(
+        model,
+        hls_config=config,
+        io_type=io_type,
+        output_dir=str(test_root_path / f'hls4mlprj_globalplool1d_{backend}_{io_type}_{model_type}_keepdims{keepdims}'),
+        backend=backend,
+    )
     hls_model.compile()
-    
-    y_keras = np.squeeze(model.predict(data_1d))
-    y_hls = hls_model.predict(data_1d)
+
+    y_keras = model.predict(data_1d)
+    y_hls = hls_model.predict(data_1d).reshape(y_keras.shape)
     np.testing.assert_allclose(y_keras, y_hls, rtol=0, atol=atol, verbose=True)
 
+
 @pytest.fixture(scope='module')
 def data_2d():
     return np.random.rand(100, in_shape, in_shape, in_filt)
 
-@pytest.fixture(scope='module')
-def keras_model_max_2d():
-    model = Sequential()
-    model.add(GlobalMaxPooling2D(input_shape=(in_shape, in_shape, in_filt)))
-    model.compile()
-    return model
 
 @pytest.fixture(scope='module')
-def keras_model_avg_2d():
+def keras_model_2d(request):
+    model_type = request.param['model_type']
+    keepdims = request.param['keepdims']
     model = Sequential()
-    model.add(GlobalAveragePooling2D(input_shape=(in_shape, in_shape, in_filt)))
+    if model_type == 'avg':
+        model.add(GlobalAveragePooling2D(input_shape=(in_shape, in_shape, in_filt), keepdims=keepdims))
+    elif model_type == 'max':
+        model.add(GlobalMaxPooling2D(input_shape=(in_shape, in_shape, in_filt), keepdims=keepdims))
     model.compile()
-    return model
+    return model, model_type, keepdims
+
 
 @pytest.mark.parametrize('backend', ['Quartus', 'Vivado'])
-@pytest.mark.parametrize('model_type', ['max', 'avg'])
+@pytest.mark.parametrize(
+    'keras_model_2d',
+    [
+        {'model_type': 'max', 'keepdims': True},
+        {'model_type': 'max', 'keepdims': False},
+        {'model_type': 'avg', 'keepdims': True},
+        {'model_type': 'avg', 'keepdims': False},
+    ],
+    ids=[
+        'model_type-max-keepdims-True',
+        'model_type-max-keepdims-False',
+        'model_type-avg-keepdims-True',
+        'model_type-avg-keepdims-False',
+    ],
+    indirect=True,
+)
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
-def test_global_pool2d(backend, keras_model_max_2d, keras_model_avg_2d, data_2d, model_type, io_type):
-    
-    if model_type == 'avg':
-        model = keras_model_avg_2d
-    else:
-        model = keras_model_max_2d
-    
+def test_global_pool2d(backend, keras_model_2d, data_2d, io_type):
+
+    model, model_type, keepdims = keras_model_2d
+
     config = hls4ml.utils.config_from_keras_model(model, default_precision='ap_fixed<32,9>', granularity='name')
 
-    hls_model = hls4ml.converters.convert_from_keras_model(model,
-                                                           hls_config=config,
-                                                           io_type=io_type,
-                                                           output_dir=str(test_root_path / f'hls4mlprj_globalplool2d_{backend}_{io_type}_{model_type}'),
-                                                           backend=backend)
+    hls_model = hls4ml.converters.convert_from_keras_model(
+        model,
+        hls_config=config,
+        io_type=io_type,
+        output_dir=str(test_root_path / f'hls4mlprj_globalplool2d_{backend}_{io_type}_{model_type}_keepdims{keepdims}'),
+        backend=backend,
+    )
     hls_model.compile()
-    
-    y_keras = np.squeeze(model.predict(data_2d))
-    y_hls = hls_model.predict(data_2d)
+
+    y_keras = model.predict(data_2d)
+    y_hls = hls_model.predict(data_2d).reshape(y_keras.shape)
     np.testing.assert_allclose(y_keras, y_hls, rtol=0, atol=atol, verbose=True)