ml-explore
Pooling layers
Proposed changes
Added MaxPool and AvgPool layers. MaxPool and AvgPool layers were requested in this issue. In this PR, I propose implementing the requested pooling operations by firstly computing sliding windows and subsequently reducing them. More precisely, we can use as_strided to compute pooling sliding windows. Then, we can simply reduce over appropriate axes to implement the desired pooling operation.
Concerns: My only concern is performance. Ideally, we should call optimized backend kernels for pooling operations. There are MPSCNNPoolingAverageNode and MPSCNNPoolingMaxNode. Similarly, there is BNNSFilterCreateLayerPooling in Accelerate. Alternatively, we could implement a kernel for window reduction.
Closes https://github.com/ml-explore/mlx/issues/25.
Checklist
Put an x in the boxes that apply.
- [x] I have read the CONTRIBUTING document
- [x] I have run
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes - [x] I have added tests that prove my fix is effective or that my feature works
- [x] I have updated the necessary documentation (if needed)
@angeloskath, do you have any performance concerns regarding pooling implementation using as_strided ?
@angeloskath I implemented the requested changes. In addition, I separated the generic implementation into MaxPooling1d, MaxPooling2d, AvgPooling1d and AvgPooling2d and updated the docs accordingly. However, I am not sure about my inheritance hierarchy.
I would also appreciate if you could take a detailed look into some of the tests. The expected results were obtained by computing outputs of torch.nn pooling layer implementations and transposing the outputs to match our channels_last convention. There could be some mistakes.
Please take another look :)
@angeloskath Could you take a look at the changes addressing your comments?
Hi guys! Any updates on this? Looks like Transformers get all the Attention (pun intended). 😅
Nearly no CNN can be implemented without pooling layers (even AlexNet, I'm trying to write VGG19 but pooling layers are needed). Can we prioritise this please???
P.S. I cannot use MLX at all until this gets merged. I am mostly into CNN.
CC: @angeloskath @awni @gboduljak
@RahulBhalley I would like to complete this PR :) I am waiting for the review.
I highly appreciate your efforts @gboduljak. I hope this gets reviewed ASAP.
@gboduljak I am coming back to this PR (sorry for taking too long). Could you perhaps put it back on top of main? Since rebase is pretty hard given the many merge commits I propose applying the diff on top of main and force pushing to the same branch. Something using the following
base=$(git merge-base main pooling-layers)
git checkout main
git checkout -b new-pooling-layers
git diff $base..pooling-layers >/tmp/pooling-layers.patch
git apply --reject /tmp/pooling-layers.patch
rm docs/src/python/nn/layers.rst.rej
git add <whatever is modified or uncommitted>
git commit -m 'Add pooling layers'
git branch -m pooling-layers old-pooling-layers
git branch -m pooling-layers
git push origin pooling-layers -f
I can do the above but then the authorship of the commit will be wrong (me instead of you).
@angeloskath Thank you for informing me. I will attempt the rebase today.
@angeloskath I did the rebase according to your instructions. Many thanks for the detailed instructions. Please take a look :)
Looks great! A few issues in the docs, please fix and then I think we can merge.
FYI some instructions on building the docs here. It's useful to build / look at them for big changes to the docs otherwise you can break them / make strange outputs.
@awni It would be great if there is a preview environment where you can browse your changes to the docs. This will eliminate the need to this manually.
It would be great if there is a preview environment where you can browse your changes to the docs. This will eliminate the need to this manually.
Interesting idea.. maybe something we can setup with CircleCI 🤔
@gboduljak actually I just pushed the docs changes myself since I already made them locally.
LGTM, @angeloskath feel free to merge when you are good with it.
Thanks for the contribution!
@angeloskath Thank you for the refactor. @awni Thank you for fixing the docs nits. I was about to push my changes, but PR was already merged :)