Support GlobalPooling2D op in TFlite

testdata/dnn/tflite/generate.py

@@ -93,77 +93,57 @@ def split(x):
 inp = np.random.standard_normal((1, 3)).astype(np.float32)
 save_tflite_model(split, inp, 'split')
 
+def keras_to_tf(model, input_shape):
+    tf_func = tf.function(
+      model.call,
+      input_signature=[tf.TensorSpec(input_shape, tf.float32)],
+    )
+    inp = np.random.standard_normal((input_shape)).astype(np.float32)
+
+    return tf_func, inp
 
 fully_connected = tf.keras.models.Sequential([
   tf.keras.layers.Dense(3),
   tf.keras.layers.ReLU(),
   tf.keras.layers.Softmax(),
 ])
 
-fully_connected = tf.function(
-      fully_connected.call,
-      input_signature=[tf.TensorSpec((1,2), tf.float32)],
-)
-
-inp = np.random.standard_normal((1, 2)).astype(np.float32)
+fully_connected, inp = keras_to_tf(fully_connected, (1, 2))
 save_tflite_model(fully_connected, inp, 'fully_connected')
 
 permutation_3d = tf.keras.models.Sequential([
-  tf.keras.layers.Permute((2,1))
+  tf.keras.layers.Permute((2, 1))
 ])
 
-permutation_3d = tf.function(
-    permutation_3d.call,
-    input_signature=[tf.TensorSpec((1,2,3), tf.float32)],
-)
-inp = np.random.standard_normal((1, 2, 3)).astype(np.float32)
+permutation_3d, inp = keras_to_tf(permutation_3d, (1, 2, 3))
 save_tflite_model(permutation_3d, inp, 'permutation_3d')
 
-# Temporarily disabled as TFLiteConverter produces a incorrect graph in this case
-#permutation_4d_0123 = tf.keras.models.Sequential([
-#  tf.keras.layers.Permute((1,2,3)),
-#  tf.keras.layers.Conv2D(3,1)
-#])
-#
-#permutation_4d_0123 = tf.function(
-#    permutation_4d_0123.call,
-#    input_signature=[tf.TensorSpec((1,2,3,4), tf.float32)],
-#)
-#inp = np.random.standard_normal((1, 2, 3, 4)).astype(np.float32)
-#save_tflite_model(permutation_4d_0123, inp, 'permutation_4d_0123')
-
-permutation_4d_0132 = tf.keras.models.Sequential([
-  tf.keras.layers.Permute((1,3,2)),
-  tf.keras.layers.Conv2D(3,1)
-])
-
-permutation_4d_0132 = tf.function(
-    permutation_4d_0132.call,
-    input_signature=[tf.TensorSpec((1,2,3,4), tf.float32)],
-)
-inp = np.random.standard_normal((1, 2, 3, 4)).astype(np.float32)
-save_tflite_model(permutation_4d_0132, inp, 'permutation_4d_0132')
-
-permutation_4d_0213 = tf.keras.models.Sequential([
-  tf.keras.layers.Permute((2,1,3)),
-  tf.keras.layers.Conv2D(3,1)
+# (1, 2, 3) is temporarily disabled as TFLiteConverter produces a incorrect graph in this case
+permutation_4d_list = [(1, 3, 2), (2, 1, 3), (2, 3, 1)]
+for perm_axis in permutation_4d_list:
+    permutation_4d_model = tf.keras.models.Sequential([
+        tf.keras.layers.Permute(perm_axis),
+        tf.keras.layers.Conv2D(3, 1)
+    ])
+
+    permutation_4d_model, inp = keras_to_tf(permutation_4d_model, (1, 2, 3, 4))
+    model_name = f"permutation_4d_0{''.join(map(str, perm_axis))}"
+    save_tflite_model(permutation_4d_model, inp, model_name)
+
+global_average_pooling_2d = tf.keras.models.Sequential([
+  tf.keras.layers.GlobalAveragePooling2D(keepdims=True),
+  tf.keras.layers.ZeroPadding2D(1),
+  tf.keras.layers.GlobalAveragePooling2D(keepdims=False)
 ])
 
-permutation_4d_0213 = tf.function(
-    permutation_4d_0213.call,
-    input_signature=[tf.TensorSpec((1,2,3,4), tf.float32)],
-)
-inp = np.random.standard_normal((1, 2, 3, 4)).astype(np.float32)
-save_tflite_model(permutation_4d_0213, inp, 'permutation_4d_0213')
+global_average_pooling_2d, inp = keras_to_tf(global_average_pooling_2d, (1, 7, 7, 5))
+save_tflite_model(global_average_pooling_2d, inp, 'global_average_pooling_2d')
 
-permutation_4d_0231 = tf.keras.models.Sequential([
-  tf.keras.layers.Permute((2,3,1)),
-  tf.keras.layers.Conv2D(3,1)
+global_max_pool = tf.keras.models.Sequential([
+  tf.keras.layers.GlobalMaxPool2D(keepdims=True),
+  tf.keras.layers.ZeroPadding2D(1),
+  tf.keras.layers.GlobalMaxPool2D(keepdims=True)
 ])
 
-permutation_4d_0231 = tf.function(
-    permutation_4d_0231.call,
-    input_signature=[tf.TensorSpec((1,2,3,4), tf.float32)],
-)
-inp = np.random.standard_normal((1, 2, 3, 4)).astype(np.float32)
-save_tflite_model(permutation_4d_0231, inp, 'permutation_4d_0231')
+global_max_pool, inp = keras_to_tf(global_max_pool, (1, 7, 7, 5))
+save_tflite_model(global_max_pool, inp, 'global_max_pooling_2d')