ml-explore / mlx

MLX: An array framework for Apple silicon
https://ml-explore.github.io/mlx/
MIT License
16.45k stars 942 forks source link

[Feature] Adaptive pooling layers #400

Open NeptuneIsTheBest opened 8 months ago

NeptuneIsTheBest commented 8 months ago

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?

gboduljak commented 8 months ago

@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 commented 8 months ago

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 commented 8 months ago

@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! :)

gboduljak commented 8 months ago

@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.

NeptuneIsTheBest commented 8 months ago

@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 commented 8 months ago

@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.

NeptuneIsTheBest commented 8 months ago

@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.

rounak commented 5 months ago

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