Tensorflow: How to Pool over Depth?

tf.nn.max_pool does not support pooling over the depth dimension which is why you get an error.

You can use a max reduction instead to achieve what you're looking for:

tf.reduce_max(input_tensor, reduction_indices=[3], keep_dims=True)

The keep_dims parameter above ensures that the rank of the tensor is preserved. This ensures that the behavior of the max reduction will be consistent with what the tf.nn.max_pool operation would do if it supported pooling over the depth dimension.

TensorFlow now supports depth-wise max pooling with tf.nn.max_pool(). For example, here is how to implement it using pooling kernel size 3, stride 3 and VALID padding:

import tensorflow as tf

output = tf.nn.max_pool(images,
                        ksize=(1, 1, 1, 3),
                        strides=(1, 1, 1, 3),
                        padding="VALID")

You can use this in a Keras model by wrapping it in a Lambda layer:

from tensorflow import keras

depth_pool = keras.layers.Lambda(
    lambda X: tf.nn.max_pool(X,
                             ksize=(1, 1, 1, 3),
                             strides=(1, 1, 1, 3),
                             padding="VALID"))

model = keras.models.Sequential([
    ..., # other layers
    depth_pool,
    ... # other layers
])

Alternatively, you can write a custom Keras layer:

class DepthMaxPool(keras.layers.Layer):
    def __init__(self, pool_size, strides=None, padding="VALID", **kwargs):
        super().__init__(**kwargs)
        if strides is None:
            strides = pool_size
        self.pool_size = pool_size
        self.strides = strides
        self.padding = padding
    def call(self, inputs):
        return tf.nn.max_pool(inputs,
                              ksize=(1, 1, 1, self.pool_size),
                              strides=(1, 1, 1, self.pool_size),
                              padding=self.padding)

You can then use it like any other layer:

model = keras.models.Sequential([
    ..., # other layers
    DepthMaxPool(3),
    ... # other layers
])