python - 如何在 PyTorch 中初始化权重？

Question

如何初始化网络的权重和偏差(例如通过 He 或 Xavier 初始化)？

Answer 1

单层

要初始化单个层的权重，请使用 torch.nn.init 中的函数.例如:

conv1 = torch.nn.Conv2d(...)
torch.nn.init.xavier_uniform(conv1.weight)

或者，您可以通过写入 conv1.weight.data(即 torch.Tensor)来修改参数。示例:

conv1.weight.data.fill_(0.01)

这同样适用于偏见:

conv1.bias.data.fill_(0.01)

nn.Sequential 或自定义nn.Module

将初始化函数传递给 torch.nn.Module.apply .它将递归地初始化整个nn.Module中的权重。

apply(fn): Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also torch-nn-init).

例子:

def init_weights(m):
    if isinstance(m, nn.Linear):
        torch.nn.init.xavier_uniform(m.weight)
        m.bias.data.fill_(0.01)

net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)