欢迎访问宙启技术站
智能推送

torch.nn.init模块中的Xavier初始化方法介绍

发布时间:2023-12-23 19:11:14

Xavier初始化,也称为Glorot初始化,是一种用于初始化神经网络权重的方法。它的目标是尽量保持前向传播和反向传播过程中信号的方差不发生太大变化,从而加速训练过程。

Xavier初始化的原理是根据输入和输出神经元数量来确定权重的初始化范围。具体而言,对于全连接层,Xavier初始化方法将权重初始化成一个随机分布,服从均匀分布或者高斯分布,具体分布方式取决于网络的激活函数。对于sigmoid或tanh激活函数,使用均匀分布;对于ReLU激活函数,使用高斯分布。

Xavier初始化的原始公式为:

w = np.random.randn(n_input, n_output) / sqrt(n_input)

其中,n_input和n_output分别代表输入和输出神经元的数量。

PyTorch中提供了torch.nn.init模块来进行Xavier初始化。其中,torch.nn.init.xavier_uniform_()方法用于均匀分布的Xavier初始化,torch.nn.init.xavier_normal_()方法用于高斯分布的Xavier初始化。两种方法的参数和返回值都相同,具体介绍如下:

torch.nn.init.xavier_uniform_(tensor, gain=1)

torch.nn.init.xavier_normal_(tensor, gain=1)

参数:

- tensor:需要初始化的张量

- gain:传递给激活函数的增益(默认为1)

返回值:

- tensor:初始化后的张量

以下是使用例子:

import torch
import torch.nn as nn
import torch.nn.init as init

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 30)
        
    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

net = Net()

# 使用xavier_uniform_方法进行初始化
init.xavier_uniform_(net.fc1.weight)
init.xavier_uniform_(net.fc2.weight)

# 使用xavier_normal_方法进行初始化
init.xavier_normal_(net.fc1.weight)
init.xavier_normal_(net.fc2.weight)

在上述例子中,定义了一个名为Net的神经网络模型,包含两个全连接层。然后使用xavier_uniform_和xavier_normal_方法对其中的权重进行初始化。根据具体的激活函数类型,可以选择不同的初始化方法。