时间戳

首页 > 解决方案 > PyTorch 中的自定义权重初始化

python - PyTorch 中的自定义权重初始化

问题描述

在 中实现custom weight initialization方法的正确方法是PyTorch什么？

我相信我不能直接向“torch.nn.init”添加任何方法，但希望用我自己的专有方法初始化模型的权重。

标签： pythonpytorch

解决方案

您可以定义一个方法来根据每一层初始化权重：

def weights_init(m):
    classname = m.__class__.__name__

    if classname.find('Conv2d') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

然后将其应用于您的网络：

model = create_your_model()
model.apply(weights_init)

推荐阅读