Skip to content

Fixes grouped Conv1d/Conv2d init to use in_channels // groups - #4544

Open
CodeWithMoin wants to merge 1 commit into
ml-explore:mainfrom
CodeWithMoin:fix-grouped-conv-init
Open

CodeWithMoin wants to merge 1 commit into
ml-explore:mainfrom
CodeWithMoin:fix-grouped-conv-init

Conversation

@CodeWithMoin

@CodeWithMoin CodeWithMoin commented Sep 22, 2026 •

Copy link
Copy Markdown

Conv1d and Conv2d computed the init bound using the full in_channels, while the weight's last dimension is in_channels // groups. For groups > 1, this made uniform init range too narrow (i.e 0.71x at 2 groups, 0.5x at 4, and 0.35x at 8). This PR fixes #4536 by using fan_in = (in_channels // groups) * kernel_size, that matches PyTorch's grouped-conv default. For the Conv3d and transposed convs, we don't have groups option, so they are unaffected. I've created a new test that checks init bound for groups 1, 2 and 4 on both layers. This test fails before the fix, and passes after the PR fix, In fact the full test_nn.py passes.

This changes only the initial random weights, so loaded checkpoints are unaffected.

Before (64 in/out channels, kernel 3, seed 0):

layer    g       weight shape fan_in   1/sqrt   max|w|  ratio
Conv1d   2        (64, 3, 32)     96   0.1021   0.0722  0.707
Conv1d   4        (64, 3, 16)     48   0.1443   0.0721  0.500
Conv2d   4     (64, 3, 3, 16)    144   0.0833   0.0417  0.500

After: ratio 1.000 for every group count, both layers.

Tests: test_nn 74 OK · test_conv 20 OK · test_conv_transpose 10 OK · test_init 11 OK

  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: I've understood the init math, checked the test results then used assistance for code and new test.

This branch has not been deployed

No deployments
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.

Grouped Conv1d / Conv2d initialization uses ungrouped fan-in

1 participant