Add support for keepdims in global pooling layers (#387)

Dobiasd · Jul 23, 2023 · e5db13c · e5db13c
1 parent c3f4c64
commit e5db13c
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 12 deletions.
diff --git a/include/fdeep/import_model.hpp b/include/fdeep/import_model.hpp
@@ -568,17 +568,19 @@ inline layer_ptr create_average_pooling_3d_layer(
 }
 
 inline layer_ptr create_global_max_pooling_3d_layer(
-    const get_param_f&, const nlohmann::json&,
+    const get_param_f&, const nlohmann::json& data,
     const std::string& name)
 {
-    return std::make_shared<global_max_pooling_3d_layer>(name);
+    const bool keepdims = data["config"]["keepdims"];
+    return std::make_shared<global_max_pooling_3d_layer>(name, keepdims);
 }
 
 inline layer_ptr create_global_average_pooling_3d_layer(
-    const get_param_f&, const nlohmann::json&,
+    const get_param_f&, const nlohmann::json& data,
     const std::string& name)
 {
-    return std::make_shared<global_average_pooling_3d_layer>(name);
+    const bool keepdims = data["config"]["keepdims"];
+    return std::make_shared<global_average_pooling_3d_layer>(name, keepdims);
 }
 
 inline layer_ptr create_upsampling_1d_layer(

diff --git a/include/fdeep/layers/global_average_pooling_3d_layer.hpp b/include/fdeep/layers/global_average_pooling_3d_layer.hpp
@@ -16,14 +16,17 @@ namespace fdeep { namespace internal
 class global_average_pooling_3d_layer : public global_pooling_layer
 {
 public:
-    explicit global_average_pooling_3d_layer(const std::string& name) :
-    global_pooling_layer(name)
+    explicit global_average_pooling_3d_layer(const std::string& name, bool keepdims) :
+    global_pooling_layer(name), keepdims_(keepdims)
     {
     }
 protected:
     tensor pool(const tensor& in) const override
     {
-        tensor out(tensor_shape(in.shape().depth_), 0);
+        const auto out_dimensions = keepdims_ ?
+            fplus::append_elem(in.shape().depth_, std::vector<std::size_t>(in.shape().rank() - 1, 1)) :
+            fplus::singleton_seq(in.shape().depth_);
+        tensor out(create_tensor_shape_from_dims(out_dimensions), 0);
         for (std::size_t z = 0; z < in.shape().depth_; ++z)
         {
             float_type val = 0;
@@ -41,6 +44,7 @@ class global_average_pooling_3d_layer : public global_pooling_layer
         }
         return out;
     }
+    bool keepdims_;
 };
 
 } } // namespace fdeep, namespace internal
diff --git a/include/fdeep/layers/global_max_pooling_3d_layer.hpp b/include/fdeep/layers/global_max_pooling_3d_layer.hpp
@@ -18,14 +18,17 @@ namespace fdeep { namespace internal
 class global_max_pooling_3d_layer : public global_pooling_layer
 {
 public:
-    explicit global_max_pooling_3d_layer(const std::string& name) :
-    global_pooling_layer(name)
+    explicit global_max_pooling_3d_layer(const std::string& name, bool keepdims) :
+    global_pooling_layer(name), keepdims_(keepdims)
     {
     }
 protected:
     tensor pool(const tensor& in) const override
     {
-        tensor out(tensor_shape(in.shape().depth_), 0);
+        const auto out_dimensions = keepdims_ ?
+            fplus::append_elem(in.shape().depth_, std::vector<std::size_t>(in.shape().rank() - 1, 1)) :
+            fplus::singleton_seq(in.shape().depth_);
+        tensor out(create_tensor_shape_from_dims(out_dimensions), 0);
         for (std::size_t z = 0; z < in.shape().depth_; ++z)
         {
             float_type val = std::numeric_limits<float_type>::lowest();
@@ -43,6 +46,7 @@ class global_max_pooling_3d_layer : public global_pooling_layer
         }
         return out;
     }
+    bool keepdims_;
 };
 
 } } // namespace fdeep, namespace internal
diff --git a/keras_export/convert_model.py b/keras_export/convert_model.py
@@ -631,8 +631,6 @@ def get_layer_weights(layer, name):
     layer_type = type(layer).__name__
     if hasattr(layer, 'data_format'):
         assert layer.data_format == 'channels_last'
-    if hasattr(layer, 'keepdims'):  # Pooling layers
-        assert not layer.keepdims
 
     show_func = get_layer_functions_dict().get(layer_type, None)
     shown_layer = None

diff --git a/keras_export/generate_test_models.py b/keras_export/generate_test_models.py
@@ -164,6 +164,8 @@ def get_test_model_exhaustive():
     outputs.append(AveragePooling1D(2, strides=2, padding='same')(inputs[6]))
     outputs.append(GlobalMaxPooling1D()(inputs[6]))
     outputs.append(GlobalAveragePooling1D()(inputs[6]))
+    outputs.append(GlobalMaxPooling1D(keepdims=True)(inputs[6]))
+    outputs.append(GlobalAveragePooling1D(keepdims=True)(inputs[6]))
 
     outputs.append(Normalization(axis=None, mean=2.1, variance=2.2)(inputs[4]))
     outputs.append(Normalization(axis=-1, mean=2.1, variance=2.2)(inputs[6]))
@@ -203,6 +205,10 @@ def get_test_model_exhaustive():
     outputs.append(GlobalAveragePooling3D()(inputs[2]))
     outputs.append(GlobalMaxPooling2D()(inputs[4]))
     outputs.append(GlobalMaxPooling3D()(inputs[2]))
+    outputs.append(GlobalAveragePooling2D(keepdims=True)(inputs[4]))
+    outputs.append(GlobalAveragePooling3D(keepdims=True)(inputs[2]))
+    outputs.append(GlobalMaxPooling2D(keepdims=True)(inputs[4]))
+    outputs.append(GlobalMaxPooling3D(keepdims=True)(inputs[2]))
 
     outputs.append(CenterCrop(4, 5)(inputs[4]))
     outputs.append(CenterCrop(5, 6)(inputs[4]))