pytorch - 如何在pytorch中查找内置函数源代码

Question

我正在尝试对批量标准化进行研究，并且不得不对 pytorch BN 代码进行一些修改。我深入研究了 pytorch 代码，并陷入了 torch.nn.functional.batch_norm ，它引用了 torch.batch_norm 。

问题是无法在 Torch 库中进一步找到 torch.batch_norm。有什么办法可以找到这个内置函数的源代码并重新实现它吗？谢谢!

Answer 1

它在那里，但它没有在 Python 中定义。它们是在 C++ 中的 aten/ 目录中定义的。

对于CPU，实现(其中之一，取决于输入是否连续)在这里:https://github.com/pytorch/pytorch/blob/420b37f3c67950ed93cd8aa7a12e673fcfc5567b/aten/src/ATen/native/Normalization.cpp#L61-L126

对于 CUDA，实现在这里:https://github.com/pytorch/pytorch/blob/7aae51cdedcbf0df5a7a8bf50a947237ac4b3ee8/aten/src/ATen/native/cudnn/BatchNorm.cpp#L52-L143