Support for `AdaptivePooling`

[awsaf49]

System information.

TensorFlow version (you are using): 2.12.0
Are you willing to contribute it (Yes/No) : Yes

Describe the feature and the current behavior/state.
Currently, Keras doesn't have support for this feature. Adaptive pooling is widely used in architecture building as it allows for dynamic kernel size and can create a desired output shape, unlike traditional pooling methods. Even PyTorch support this AdaptivePooling.

Use-case: Let's consider an image classification task where the network needs to accept images of varying sizes. We want the network to output a fixed-size feature map regardless of the input image size. Here, adaptive pooling would be useful. By dynamically adjusting the pooling window size and stride based on the input, adaptive pooling ensures that the output feature map size is constant.

Will this change the current api? How?
No

Who will benefit from this feature?
The research community who are working on creating new architecture will benefit from this.

Contributing

• Do you want to contribute a PR? (yes/no): Yes
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):
  As TensorFlow Addons has support for this feature here, I would like to adapt it from there with the necessary modifications to fit Keras API. This feature may look something like this,

tf.keras.layers.AdaptiveAveragePooling2D(output_size, data_format=None, **kwargs)

[innat]

keras-team/keras-cv#512
cc @ianstenbit to reopen ticket on kcv. @Rocketknight1

[Rocketknight1]

lol, hi again @awsaf49! Great minds think alike.

I investigated this before (in July 2022) and the version in Tensorflow Addons did not match the Torch behaviour, which meant it was not usable for porting Torch models. I wrote a correct port that always matches Torch here. However, I'm not sure there was enough interest from the Keras-CV team at the time. We just added my implementation to the relevant models in transformers.

[ianstenbit]

I am open to re-opening that issue on KerasCV if @Rocketknight1 is interested in contributing this layer to KerasCV.

Most likely this layer doesn't belong here (in core Keras), unless there are some non-vision use cases for it. That said, I see that it's relevant for some vision transformers, so we should include it in KerasCV

[Rocketknight1]

I'm open to submitting a PR to KerasCV for it! It's unclear what the right approach is, though. For us at Hugging Face, the most important thing is that it matches the behaviour of PyTorch. This is because a lot of big models are trained in PyTorch, so if we want their pretrained weights to work in TensorFlow without retraining then we need a layer whose behaviour is exactly the same.

However, the way the PyTorch layer works is quite strange: it tiles a lot of non-overlapping windows with variable width across the entire image. This makes it very unnatural to implement without their custom kernel (you can see how weird my custom layer had to be above). A "natural" TF AdaptivePool would probably do things differently, but that would lose the PyTorch compatibility.

[ianstenbit]

I would say our highest priority is an elegant API and straightforward implementation, and matching the PyTorch version is more of a nice-to-have. So I'd say build it how it naturally would exist in TF, and in our docstring we can point users to your PyTorch-mirrored version for those folks that are looking to port models.

How does that sound to you?

[Rocketknight1]

That could work! How would you feel about a match_torch_behavior argument to the constructor, with the layer able to flip between both modes, or is that too much complexity?

[awsaf49]

Most likely this layer doesn't belong here (in core Keras), unless there are some non-vision use cases for it.

@ianstenbit AdaptivePooling can be applied to 1D features using AdaptivePooling1D layer thus I created this issue here instead of keras-cv. So, not sure if AdaptivePooling is limited to vision applications only

[awsaf49]

Tensorflow Addons did not match the Torch behaviour, which meant it was not usable for porting Torch models.

@Rocketknight1, WoW, I wasn't aware of this mismatch. I tried it myself, and you're correct. In TensorFlow, it splits the input dimensions into each output dimension, which means the input dimension needs to be divisible by the input shape. As a result, it can't handle any input shape for any output shape. Also, your code for AdaptivePooling seems to behave similarly to pytorch version.
Here is a demo notebook.

In that context, I think a version which is compatible with PyTorch would be better for porting or implementing paper.

[ianstenbit]

From what I can tell, this is predominantly used in vision applications today, so let's start out by adding this to KerasCV -- please feel free to open an issue and/or PR on KerasCV.

If we later determine that this has broad enough usage to upstream to Keras, that's always an option. I'm closing this issue here, but let's move the conversation to a new issue on KerasCV 😄

Thanks!