Pytorch深入浅出（五）之权值初始化

原创于 2025-12-19 08:00:00 发布 · 1.2k 阅读

11 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习

Pytoch 专栏收录该内容

9 篇文章

订阅专栏

在搭建好网络模型后，训练的第一步不是输入数据，而是初始化权值。如果把训练神经网络比作“滚雪球”，初始化就是决定雪球从山顶哪里开始滚。位置没选好，雪球可能根本滚不动（梯度消失），或者瞬间滚散架（梯度爆炸）。

引言：初始化的终极目标

“保持方差一致性” —— 让数据在经过每一层网络时，信号的强度（方差）保持不变。

前向传播时：每一层输出的方差 $\approx$ 输入的方差（信号不衰减、不放大）。
反向传播时：每一层梯度的方差 $\approx$ 上一层梯度的方差（梯度流得动）。

💡 黄金法则：

只初始化“带可训练参数”的层
Weight (权重)：决定信号的“放大倍数”，是控制方差的核心。必须根据激活函数的特性精心选择初始化方法（如 ReLU 配合 Kaiming）。
Bias (偏置)：只负责平移，不影响方差。通常一律初始化为 0 即可。

初始化是为了让梯度能在网络里‘流得动’。只要梯度流得动（看 Tensorboard 的梯度直方图），标准规则就是最好的规则。

一、预备知识：增益 (Gain) 与 calculate_gain

不同的激活函数对信号方差的影响不同。例如，Tanh 会把数据压缩到 (-1, 1)，导致方差变小；ReLU 会把一半的数据变为 0，也会改变方差。
为了抵消这种影响，PyTorch 引入了 “增益 (Gain)” 的概念。它是一个补偿系数，用来把变小的方差“拉回来”。

torch.nn.init.calculate_gain(_nonlinearity_, _param=None_)
“““
@params：
	nonlinearity：非线性函数（nn.functional 名称）['linear', 'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d', 'sigmoid', 'tanh', 'relu', 'leaky_relu', 'selu']
	param：非线性函数的可选参数
@return
	返回给定非线性函数的推荐增益值
”””

常见激活函数的推荐增益：

非线性激活函数	推荐 Gain 值	直观理解
Linear / Identity	$1$	线性变换，不改变方差规模
Sigmoid	$1$	(主要用于 Xavier 初始化，不做额外增益)
Tanh	$5/3$	Tanh 会压缩数据，需要放大权重来补偿
ReLU	$\sqrt{2}$	ReLU 砍掉了一半数据（负半轴），需要 $\sqrt{2}$ 倍增益来维持方差
Leaky ReLU	$\sqrt{2 / (1 + \alpha^2)}$	根据负轴斜率 $\alpha$ 动态调整

二、初始化方法

PyTorch 的 torch.nn.init 模块提供了多种方法，在此模块中的所有函数都旨在用于初始化神经网络参数，因此它们都在 torch.no_grad() 模式下运行，并且不会被 autograd 考虑在内。我们可以将其分为三个阶段来理解：

1.阶段一：朴素初始化

均匀分布 (uniform_)：从 $U (a, b)$ 中采样。
正态分布 (normal_)：从 $N(\text{mean}, \text{std}^2)$ 中采样。
常数分布 (constant_)：用固定值填充（常用于 bias 初始化为 0）。

具体函数：

torch.nn.init.uniform_(tensor, a=0.0, b=1.0, generator=None)

功能：将输入张量的值用均匀分布 $U (a, b)$ 随机采样得到的值填充。

tensor：要初始化的张量
a、b：均匀分布的上下界

torch.nn.init.normal_(tensor, mean=0.0, std=1.0, generator=None)

功能：将输入张量的值用正态分布 $N(mean,std ^2)$ 随机采样得到的值填充。

tensor：要初始化的张量
mean：正态分布的均值
std：正态分布的标准差

torch.nn.init.constant_(tensor, val)

用固定值 val 去填充输入张量。

tensor：要填充的张量
val：要填充的值

2.阶段二：Xavier / Glorot 初始化 (Sigmoid/Tanh 时代)
适用场景： 激活函数为 Sigmoid 或 Tanh 的网络。
在 2010 年由 Glorot 和 Bengio 提出。为了解决深层网络中 Sigmoid 导致的梯度消失问题。它根据输入和输出神经元的数量（fan_in, fan_out）自动调整方差。

Xavier 均匀分布 (xavier_uniform_)：
$\sim U \left[ -\text{gain} \times \sqrt{\frac{6}{n_{in} + n_{out}}}, \quad \text{gain} \times \sqrt{\frac{6}{n_{in} + n_{out}}} \right]$
Xavier 正态分布 (xavier_normal_)：
$\sim N \left( 0, \quad \left(\text{gain} \times \sqrt{\frac{2}{n_{in} + n_{out}}}\right)^2 \right)$
具体函数：

torch.nn.init.xavier_uniform_(tensor, gain=1.0, generator=None)

功能：使用 Xavier 均匀分布填充输入 Tensor 的值。
$\sim U\left[ -\frac{\sqrt{6}}{\sqrt{n_i + n_{i+1}}}, \frac{\sqrt{6}}{\sqrt{n_i + n_{i+1}}} \right]$
生成的张量将从中采样值 $U (- a, a)$ ，其中 $\text{gain} \times \sqrt{\frac{6}{fan_{in} + fan_{out}}}$

tensor：要初始化的张量
gain：可选的缩放因子。根据激活函数来定，保证网络层各层权重的方差差距不大

torch.nn.init.xavier_normal_(tensor, gain=1.0, generator=None)

功能：使用 Xavier 正态分布填充输入 Tensor 的值。
生成的张量将从中采样值 $N(mean,std ^2)$ ，其中 $\text{gain} \times \sqrt{\frac{2}{fan_{in} + fan_{out}}}$

tensor：要初始化的张量
gain：可选的缩放因子

3.阶段三：Kaiming / He 初始化 (ReLU 时代)✨
适用场景： 激活函数为 ReLU 或 Leaky ReLU 的网络。
在 2015 年由 Kaiming He（何恺明）提出。因为 Xavier 假设激活函数是线性的（在零点附近），但 ReLU 砍掉了一半激活值，导致 Xavier 初始化的网络方差会逐层减半。Kaiming 初始化在 Xavier 的基础上乘了一个 $\sqrt{2}$ 的系数。

Kaiming 均匀分布 (kaiming_uniform_):
$fan_mode 边界：\text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}}$
Kaiming 正态分布 (kaiming_normal_):
$fan_mode 标准差：\text{std} = \frac{\text{gain}}{\sqrt{\text{fan\_mode}}}$

具体函数：

torch.nn.init.kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu', generator=None)

功能：使用 Kaiming 均匀分布填充输入 Tensor 的值。
生成的张量将从中采样值 $U (- b o u n d, b o u n d)$ ，其中 $fan_mode \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}}$

torch.nn.init.kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu', generator=None)

使用 Kaiming 正态分布填充输入 Tensor 的值。
生成的张量将从中采样值 $N(mean,std ^2)$ ，其中 $\mathrm{std} = \frac{\mathrm{gain}}{\sqrt{\mathrm{fan\_mode}}}$

tensor：要初始化的张量
a：此层之后使用的整流器的负斜率（仅当使用 'leaky_relu' 时使用）
mode：'fan_in'（默认）或 'fan_out'。选择 'fan_in' 可保持前向传播中权值方差的大小。选择 'fan_out' 可保持反向传播中的大小。
nonlinearity：非线性函数（nn.functional 名称），建议仅与 'relu' 或 'leaky_relu'（默认）一起使用。

三、其他高级初始化

截断正态分布 (trunc_normal_)：从正态分布采样，但切除掉 2 个标准差之外的离群点。Transformer (如 BERT) 模型中常用，比普通正态分布更稳定。
正交初始化 (orthogonal_)：用正交矩阵初始化权重。常用于 RNN / LSTM 等循环神经网络，能有效防止梯度消失/爆炸。
稀疏初始化 (sparse_)：将矩阵初始化为稀疏矩阵，大部分元素为 0。

具体函数：

torch.nn.init.trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0, generator=None)

功能：使用截断正态分布填充输入张量的值。
这些值实际上是从正态分布 $N(mean,std ^2)$ 中抽取的，并在超出 $[a, b]$ 范围的值被重新绘制直到它们在范围内。生成随机值的方法在 $a \leq m e an \leq b$ 时效果最好。

tensor：要初始化的张量
mean：正态分布的均值
std：正态分布的标准差
a、b：最小/大截断值

torch.nn.init.orthogonal_(tensor, gain=1, generator=None)

功能：用（半）正交矩阵填充输入 Tensor。
输入张量必须至少有 2 个维度，对于大于 2 个维度的张量，其后面的维度将被展平。

tensor：一个 n 维 torch.Tensor，其中 n≥2
gain：可选的缩放因子

torch.nn.init.sparse_(_tensor_, _sparsity_, _std=0.01_, _generator=None_)

功能：将二维输入 Tensor 填充为稀疏矩阵。
非零元素将从正态分布 $N (0, 0.01)$ 中抽取

tensor：一个 n 维 torch.Tensor
sparsity：每列中设置为零的元素的比例
std：用于生成非零值的正态分布的标准差

四、代码实践

最主要是定义一个函数weights_init ，对“有可训练参数”的层，根据激活函数选择初始化方法

import os
import torch
import random
import numpy as np
import torch.nn as nn
import matplotlib.pyplot as plt
torch.manual_seed(42)

class MyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(32 * 20 * 20, 10)
  
    def forward(self, x):
        x = self.conv1(x)
        print(f"conv1 std: {x.std():.6f}")
        x = self.relu(x)
        print(f"relu  std: {x.std():.6f}")
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        print(f"fc    std: {x.std():.6f}")
        return x

    def weights_init(self):
        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.Linear)):
        # 【重点】根据激活函数选择方法
        # 假设我们用的是 ReLU，所以使用 kaiming_normal
                nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
        # Bias 一律初始化为 0
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
    # 如果是 BatchNorm 层 (通常 weight为1, bias为0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

net = MyNet()
net.weights_init()

pic = torch.randn(16, 3, 22, 22)   # 任意给一张小图
with torch.no_grad():
    out = net(pic)

输出每层的std

conv1 std: 1.495022 
relu std: 0.875383 
fc std: 1.382164