DOC Adds code comment for _ConvNd.reset_parameters (#58931)

thomasjpfan · facebook-github-bot · commit 8130f2f67ada · 2021-05-26T08:39:40.000-07:00
Summary: Fixes #55741 by adding a comment regarding the behavior of `kaiming_uniform_` The docstring is correct in this case. For example: ```python import math import matplotlib.pyplot as plt import torch import torch.nn as nn in_channels = 120 groups = 2 kernel = (3, 8) m = nn.Conv2d(in_channels=in_channels, groups=groups, out_channels=100, kernel_size=kernel) k = math.sqrt(groups / (in_channels * math.prod(kernel))) print(f"k: {k:0.6f}") print(f"min weight: {m.weight.min().item():0.6f}") print(f"max weight: {m.weight.max().item():0.6f}") ``` outputs: ``` k: 0.026352 min weight: -0.026352 max weight: 0.026352 ``` And when we plot the distribution, it is uniform with the correct bounds: ```python _ = plt.hist(m.weight.detach().numpy().ravel()) ``` ![Unknown](https://user-images.githubusercontent.com/5402633/119552979-21ba3800-bd69-11eb-8e10-e067c943abe3.png) Pull Request resolved: #58931 Reviewed By: anjali411 Differential Revision: D28689863 Pulled By: jbschlosser fbshipit-source-id: 98eebf265dfdaceed91f1991fc4b1592c0b3cf37
diff --git a/torch/nn/modules/conv.py b/torch/nn/modules/conv.py
@@ -138,6 +138,9 @@ def __init__(self,
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
+        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
+        # uniform(-1/sqrt(k), 1/sqrt(k)), where k = weight.size(1) * prod(*kernel_size)
+        # For more details see: https://github.com/pytorch/pytorch/issues/15314#issuecomment-477448573
         init.kaiming_uniform_(self.weight, a=math.sqrt(5))
         if self.bias is not None:
             fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
