Mobile wallpaper 1Mobile wallpaper 2Mobile wallpaper 3Mobile wallpaper 4
2176 字
11 分钟
深度学习计算:延后初始化
2026-03-22
统计加载中...

本文内容参考 Dive into Deep Learning 1st Edition by Aston Zhang (Author), Zachary C. Lipton (Author), Mu Li (Author), Alexander J. Smola (Author) 在构建深度学习模型时,我们常常会遇到一个看似简单却容易出错的问题:如何确定每一层神经网络的输入维度? 尤其是在设计复杂的网络架构时,手动计算和传递这些维度不仅繁琐,还极易出错。一个微小的计算失误就可能导致整个程序崩溃。

你是否也遇到过这样的情况?

  • 定义了一个网络框架,但无法立即确定输入数据的特征维度。
  • 添加网络层时,不清楚前一层的输出维度是多少。
  • 在初始化模型参数时,由于信息不足,无法确定权重矩阵的形状。

传统的解决方案要求我们在定义网络时就必须精确指定每一层的输入和输出维度。然而,PyTorch提供了一种更优雅、更灵活的机制来解决这个问题——延后初始化(Lazy Initialization)

什么是延后初始化?#

简单来说,延后初始化允许我们“先搭框架,后定参数”。当我们使用如 nn.LazyLinear 这样的“懒惰”层时,我们只需要指定其输出维度。至于输入维度,框架会耐心等待,直到第一次有真实数据通过网络时,才根据数据的形状动态推断出来,并完成参数的创建和初始化。

你可以把 nn.LazyLinear 看作是 nn.Linear 的“拖延症”版本,但这种拖延是极具智慧的。

特性nn.Linear (普通)nn.LazyLinear (懒惰)
初始化时机定义时立即创建参数第一次执行前向传播时创建
必需参数in_featuresout_features仅需 out_features
容错性维度算错则程序立即报错自动适配输入,不易出错
内存占用定义后立即占用参数内存数据传入前几乎不占内存

实践:从代码看延后初始化#

让我们通过一个具体的例子,一步步揭开延后初始化的神秘面纱。

第一步:搭建一个“懒惰”的网络#

import torch
from torch import nn
# 使用 LazyLinear 构建一个简单的两层网络
# 我们只关心每一层要输出多少特征,而不关心输入特征是多少
net = nn.Sequential(
nn.LazyLinear(256), # 输出256维,输入?未知!
nn.ReLU(),
nn.LazyLinear(10) # 输出10维(例如10分类),输入?也未知!
)
print("定义网络后,第一层的权重是:", net[0].weight)
# 输出:None

此时,net 只是一个“空架子”。两个 LazyLinear 层记录了自己的输出维度目标(256和10),但它们的权重(weight)和偏置(bias)参数都还是 None,因为它们不知道输入维度,无法创建具体大小的张量。

第二步:用数据触发初始化#

现在,我们制造一些模拟数据,并让它流过网络。

# 假设我们的数据批次大小为2,每个样本有20个特征
X = torch.rand(2, 20)
# 第一次前向传播!魔法在此发生
output = net(X)
print("\n第一次前向传播后...")
print("第一层的权重形状:", net[0].weight.shape)
print("第二层的权重形状:", net[1].weight.shape)

运行这段代码,你会看到类似以下的输出:

第一层的权重形状:torch.Size([256, 20])
第二层的权重形状:torch.Size([10, 256])

发生了什么?

  1. 数据 X 进入第一层 LazyLinear(256):框架检测到 X 的形状为 (2, 20),立即推断出 in_features=20。于是,它瞬间将 LazyLinear(256) “变身”为一个真正的 Linear(20, 256) 层,并创建了形状为 (256, 20) 的权重和 (256,) 的偏置。
  2. 数据继续流动:第一层输出形状为 (2, 256) 的张量,经过ReLU激活函数后,进入第二层 LazyLinear(10)
  3. 第二层初始化:框架看到输入特征为256,于是将 LazyLinear(10) 实例化为 Linear(256, 10),创建了形状为 (10, 256) 的权重和 (10,) 的偏置。

至此,我们的“懒惰”网络 net 已经完成了蜕变,内部结构变成了: Linear(20, 256) -> ReLU -> Linear(256, 10)

整个过程完全自动化,我们无需手动计算和传递任何输入维度。

深入原理:初始化流程与自定义控制#

延后初始化解决了“何时创建参数”的问题,但“如何初始化这些参数”同样重要。PyTorch的默认初始化策略可能并不总是最优的。因此,我们常常需要在参数创建后,立即对其应用特定的初始化方法(如Xavier或Kaiming初始化)。

下面是一个常见的工具函数,它完美地结合了延后初始化和自定义初始化:

def apply_init(self, inputs, init=None):
# 步骤1:执行前向传播,触发延后初始化
self.forward(*inputs)
# 步骤2:如果提供了初始化函数,则应用到网络的所有参数上
if init is not None:
self.net.apply(init)

让我们拆解这个函数,理解每一行代码的“语法角色”和“实际作用”:

代码片段语法角色实际传入的内容通俗比喻
self实例引用模型对象本身“我”自己
inputs位置参数一个包含数据的元组,如 (X,)一箱待处理的零件
*inputs参数解包元组里的具体张量数据 X把零件从箱子里拿出来用
init关键字参数某种初始化算法(函数对象)一桶指定颜色的油漆

工作流程详解#

假设我们这样调用它:apply_init(net, (X,), init=my_init_func)

  1. self.forward(*inputs)

    • 这是整个流程的关键触发点*inputs 将元组 (X,) 解包,变成 forward(X)
    • 这行代码的唯一目的,就是让数据 X 流经一次网络。如前所述,这个动作会迫使所有 LazyLinear 层根据 X 的形状推断出自己的输入维度,并创建出实实在在的参数张量(weightbias)。
    • 重要澄清:这一步只创建参数,不改变参数的值。参数的值仍然是PyTorch的默认初始值(如均匀分布)。它也不会改变 init 变量本身,init 是否为 None 完全由调用者决定。
  2. if init is not None:

    • 这是一个条件检查。它判断用户是否传入了自定义的初始化函数 init
    • 如果 initNone,函数到此结束,参数保持默认初始化状态。
    • 如果 init 是一个有效的函数,则执行下一步。
  3. self.net.apply(init)

    • applynn.Module 的一个方法,它会递归地将传入的函数 init 应用到网络 self.net 中的每一个子模块上。
    • 此时,由于第一步已经执行,网络中所有层的参数都已经存在。init 函数就可以安全地访问并修改这些参数。
    • 一个典型的 init 函数长这样:
    def my_init_func(module):
    if isinstance(module, nn.Linear):
    # 对线性层的权重使用Xavier均匀初始化
    nn.init.xavier_uniform_(module.weight)
    # 将偏置项初始化为0
    if module.bias is not None:
    nn.init.zeros_(module.bias)
    • apply 遍历到 net[0](即第一个 Linear 层)时,会调用 my_init_func(net[0]),从而将其 weight 重新初始化为Xavier分布,bias 置零。

总结与最佳实践#

延后初始化是一种强大的设计模式,它极大地简化了深度学习模型的构建过程,特别是在研究原型设计和快速迭代阶段。它的核心价值在于:

  • 降低复杂度:开发者无需手动追踪和计算层与层之间的维度变化。
  • 提升容错性:避免了因手动输入维度错误导致的运行时崩溃。
  • 灵活高效:在定义复杂或条件网络时(如输入维度可配置),代码更加清晰。同时,参数内存的占用被推迟,在模型定义阶段更节省资源。

在使用时,请记住以下要点:

  • 仅第一次前向传播特殊:只有第一次调用 forward 时会触发初始化。之后的所有调用,网络都像普通网络一样工作。
  • 初始化顺序很重要:务必先通过一次前向传播(或调用类似 apply_init 的函数)完成参数创建,然后再进行自定义的参数初始化。
  • 并非所有层都支持:目前主要是 LazyLinearLazyConv2d 等基础层支持延后初始化。对于自定义层,需要实现相应的逻辑才能支持。

通过掌握延后初始化,你可以将更多精力专注于模型结构的设计和实验,而将繁琐的维度管理交给框架,让代码既简洁又健壮。

深度学习计算:延后初始化
https://blog.solmount.top/posts/lazy-init/
作者
空 柏
发布于
2026-03-22
许可协议
CC BY-NC-SA 4.0

部分信息可能已经过时

封面
Sample Song
Sample Artist
封面
Sample Song
Sample Artist
0:00 / 0:00