Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add groups to 2-D convolutions #1129

Merged
merged 2 commits into from
May 23, 2024
Merged

Add groups to 2-D convolutions #1129

merged 2 commits into from
May 23, 2024

Conversation

Rifur13
Copy link
Contributor

@Rifur13 Rifur13 commented May 16, 2024

Proposed changes

Added groups to 2-D convolutions for some kernel specializations.

Also fixed 1D grouped convs with different kernel strides and added more tests.

Can close out #100

Performance looks pretty good:

(N, H, W, C) (O, kH, kW, C) dtype stride pads groups diff%
(4, 64, 64, 256) (256, 5, 5, 256) float32 (1, 1) (2, 2) 1 +25.78%
(4, 64, 64, 256) (256, 5, 5, 256) float32 (1, 1) (2, 2) 2 +36.72%
(4, 64, 64, 256) (256, 5, 5, 256) float32 (1, 1) (2, 2) 16 -23.80%
(4, 64, 64, 256) (256, 5, 5, 256) float32 (1, 1) (2, 2) 64 +92.72%

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

…alizations.

Also fixed 1D grouped convs with different kernel strides and added more tests.
Copy link
Member

@awni awni left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice results! 🚀

Code looks good to me as well!

@jagrit06 do you want to take a look before I merge?

Copy link
Member

@jagrit06 jagrit06 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me for now! Few things can be updated here and there in the metal kernel, but I can do that in a later optimization pass

Thanks for the great work @Rifur13 !

@awni awni merged commit 9401507 into ml-explore:main May 23, 2024
5 checks passed
jkaercher pushed a commit to jkaercher/mlx that referenced this pull request May 30, 2024
* Added groups to 2-D convolutions. Only implemented for **some** specializations.

Also fixed 1D grouped convs with different kernel strides and added more tests.

* fix channels condition
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants