torch module.apply(fn)说明

蓝羽飞鸟

已于 2022-03-22 10:26:33 修改

阅读量943

点赞数 1

CC 4.0 BY-SA版权

分类专栏： pytorch 文章标签：深度学习计算机视觉

于 2022-03-22 10:21:20 首次发布

本文链接：https://blog.youkuaiyun.com/level_code/article/details/123653677

pytorch 专栏收录该内容

12 篇文章

订阅专栏

本文介绍了如何在PyTorch中使用nn.Module及其子类，并详细讲解了如何定义weight_init函数来初始化子模块的权重，以及如何通过apply()方法递归地应用该函数。通过实例展示，读者将学会如何正确地对神经网络模型进行权重初始化。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

应用的对象应是nn.Module类型或者它的继承类。
可以有children。

比如

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))
net = Model()

这里的net就是一个Module对象，它包含了children: conv1, conv2

fn用来给childer初始化weight,
比如fn可以这样定义：

def weight_init(m):
    """ Weight initialization function. """
    # Conv2D
    if isinstance(m, nn.Conv2d):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    else:
        pass

用的时候Applies fn recursively to every submodule (as returned by .children()) as well as self.
可看到官方源码如下，先对children应用fn, 再对self应用，最后返回self

def apply(self: T, fn: Callable[['Module'], None]) -> T:
     for module in self.children():
         module.apply(fn)
     fn(self)
     return self

应用例：

net.apply(weight_init)