如何在PyTorch中初始化权重? [英] How to initialize weights in PyTorch?

问题描述

如何在PyTorch的网络中初始化权重和偏差(例如，使用He或Xavier初始化)?

How to initialize the weights and biases (for example, with He or Xavier initialization) in a network in PyTorch?

推荐答案

单层

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

Single layer

To initialize the weights of a single layer, use a function from torch.nn.init. For instance:

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

或者，您可以通过写入conv1.weight.data(这是 torch.Tensor ).示例:

Alternatively, you can modify the parameters by writing to conv1.weight.data (which is a torch.Tensor). Example:

conv1.weight.data.fill_(0.01)

偏见也是如此:

conv1.bias.data.fill_(0.01)

nn.Sequential或自定义nn.Module

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

nn.Sequential or custom nn.Module

Pass an initialization function to torch.nn.Module.apply. It will initialize the weights in the entire nn.Module recursively.

apply( fn )::递归地将fn应用于每个子模块(由.children()返回)和自身.典型的用法包括初始化模型的参数(另请参见torch-nn-init).

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 type(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)

这篇关于如何在PyTorch中初始化权重?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！