mlx
mlx copied to clipboard
[Feature] Adaptive pooling layers
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 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?
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
.
@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
andmx.mean
.
I will try to implement an AdaptiveAvgPooling myself first, thanks for your suggestions! Hope to see these Pooling officially released soon! :)
@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
andmx.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 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? :>
@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 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.
As an ML beginner, would love it if there was support for AdaptiveAvgPool2d.