mlx icon indicating copy to clipboard operation
mlx copied to clipboard

[Feature] Adaptive pooling layers

Open NeptuneIsTheBest opened this issue 1 year ago • 8 comments

Hello, when I looked at the PR #357 I found that there are only a few basic pooling layer implementations such as MaxPooling and AvgPooling. Will we add an adaptive pooling layer implementation similar to that in PyTorch?

NeptuneIsTheBest avatar Jan 08 '24 03:01 NeptuneIsTheBest

@NeptuneIsTheBest This is a good suggestion. I am willing to implement this after we get some feedback on MaxPooling and AvgPooling. I would like to add these in a separate PR. @awni @angeloskath should we address this within the existing PR https://github.com/ml-explore/mlx/pull/357?

gboduljak avatar Jan 08 '24 10:01 gboduljak

Hello, when I looked at the PR #357 I found that there are only a few basic pooling layer implementations such as MaxPooling and AvgPooling. Will we add an adaptive pooling layer implementation similar to that in PyTorch?

Meanwhile, if you need this feature earlier, you can implement these easily using mx.max and mx.mean.

gboduljak avatar Jan 08 '24 10:01 gboduljak

@NeptuneIsTheBest This is a good suggestion. I am willing to implement this after we get some feedback on MaxPooling and AvgPooling. I would like to add these in a separate PR. @awni @angeloskath should we address this within the existing PR #357?

I'm looking to port my model (based on MobileNet-v3-Small) from PyTorch to mlX, so I'm looking to see if there are plans for AdaptiveAvgPooling.

Hello, when I looked at the PR #357 I found that there are only a few basic pooling layer implementations such as MaxPooling and AvgPooling. Will we add an adaptive pooling layer implementation similar to that in PyTorch?

Meanwhile, if you need this feature earlier, you can implement these easily using mx.max and mx.mean.

I will try to implement an AdaptiveAvgPooling myself first, thanks for your suggestions! Hope to see these Pooling officially released soon! :)

NeptuneIsTheBest avatar Jan 08 '24 11:01 NeptuneIsTheBest

@NeptuneIsTheBest This is a good suggestion. I am willing to implement this after we get some feedback on MaxPooling and AvgPooling. I would like to add these in a separate PR. @awni @angeloskath should we address this within the existing PR #357?

I'm looking to port my model (based on MobileNet-v3-Small) from PyTorch to mlX, so I'm looking to see if there are plans for AdaptiveAvgPooling.

Hello, when I looked at the PR #357 I found that there are only a few basic pooling layer implementations such as MaxPooling and AvgPooling. Will we add an adaptive pooling layer implementation similar to that in PyTorch?

Meanwhile, if you need this feature earlier, you can implement these easily using mx.max and mx.mean.

I will try to implement an AdaptiveAvgPooling myself first, thanks for your suggestions! Hope to see these Pooling officially released soon! :)

@NeptuneIsTheBest In this implementation of MobileNetV3, adaptive pooling is used to implement global average pooling. This can be implemented without adaptive pooling, namely using mlx.mean(x, (1,2)). The calculation assumes MLX channels-last convention (e.g. input is of the shape (B, H, W, C)). For this, you do not need a full implementation of adaptive pooling.

gboduljak avatar Jan 08 '24 17:01 gboduljak

@gboduljak You are right, but I actually use AdaptiveAvgPool2D and AdaptiveMaxPool2D in many places. And (as a consequence of not spending the money to upgrade to 36GB of unified memory) I have to constantly tweak the model to test whether I'm going to run into those memory issues (I've added too many things to it, like attention and feature fusion, etc.). From an ease of use and code readability perspective, maybe implementing one in its entirety would be a better option? :>

NeptuneIsTheBest avatar Jan 08 '24 21:01 NeptuneIsTheBest

@gboduljak You are right, but I actually use AdaptiveAvgPool2D and AdaptiveMaxPool2D in many places. And (as a consequence of not spending the money to upgrade to 36GB of unified memory) I have to constantly tweak the model to test whether I'm going to run into those memory issues (I've added too many things to it, like attention and feature fusion, etc.). From an ease of use and code readability perspective, maybe implementing one in its entirety would be a better option? :>

I think having a full implementation of all 'adaptive' pooling layers is useful. It might be worthy to wait until 'simple' pooling layers are implemented, because it is possible to use these to implement the 'adaptive' ones. My comment was only to give a quick temporary solution, until this is done.

gboduljak avatar Jan 08 '24 22:01 gboduljak

@gboduljak I've been implementing this as a temporary solution based on the method mentioned in #25 and provided by you. :>This should get me through the period leading up to the official release.

NeptuneIsTheBest avatar Jan 09 '24 19:01 NeptuneIsTheBest

As an ML beginner, would love it if there was support for AdaptiveAvgPool2d.

rounak avatar Apr 06 '24 21:04 rounak