本文基于d2l项目内容整理，介绍深度学习中的延后初始化技术，包括其原理、应用场景和实现方法。

1. 延后初始化的概念

1.1 传统初始化 vs 延后初始化

传统初始化特点：

  • 模型参数在模型创建时被立即初始化
  • 需要预先知道所有层的输入维度
  • 参数在模型定义阶段就确定大小
  • 内存在模型创建时就被分配

延后初始化特点：

  • 推迟参数初始化的时机到首次传入数据后
  • 根据实际输入动态推断各个层的参数大小
  • 在输入维度未知的情况下预定义灵活的模型
  • 延迟内存分配，提高资源利用效率

通常情况下，模型参数在模型创建时被立即初始化。但有时，使用延后初始化 (deferred initialization) 技术能推迟参数初始化的时机，直到首次传入数据后，才初始化参数。这种技术旨在输入维度未知的情况下，预定义灵活的模型，动态推断各个层的参数大小。

1.2 设计模式类比

延后初始化与软件设计模式中的懒汉模式 (Lazy Singleton) 有相似的延迟操作思想：

1
2
3
4
5
6
7
class Singleton:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

设计思想相似性：

懒汉模式是实现单例的方式之一，保证类在程序运行期间仅有一个实例，该实例仅在第一次被访问时才创建。这与延后初始化都能避免提前的内存占用和对象创建。

2. 应用场景

2.1 主要应用场景

延后初始化通常在以下场景中使用：

输入数据形状不确定：

  • 处理可变长度序列数据
  • 图像尺寸不固定的计算机视觉任务
  • 批处理大小动态变化的场景
  • 多模态数据融合时的维度适配

延后初始化允许模型根据实际输入动态调整结构，无需提前指定参数形状。

大型模型内存优化：

  • 包含大量参数的大型模型
  • 避免在模型构建时占用过多内存资源
  • 直到真正访问参数时才动态分配内存
  • 提高内存使用效率

对于包含大量参数的大型模型，延后初始化可以避免在模型构建时占用过多内存资源，直到真正访问这些参数时才动态分配。

提高模型灵活性：

  • 构建通用的模型架构
  • 支持不同输入维度的复用
  • 简化模型定义过程
  • 减少硬编码的维度参数

2.2 注意事项

潜在问题：

由于延后初始化，一些潜在的参数配置问题可能要到实际调用时才被暴露出来，而不是模型构建时就可检测到，这增加了调试的难度。

建议在开发阶段进行充分的测试，确保各种输入情况下的参数初始化都能正常工作。

3. 实现方法

3.1 手动实现

基本实现思路

我们可以手动实现延后初始化的逻辑：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import torch
import torch.nn as nn


class LazyLinear(nn.Module):
    def __init__(self, out_features=10):
        super().__init__()
        self.out_features = out_features
        self.linear = None  # 不在初始化时设置 Linear 层

    def forward(self, data):
        if self.linear is None:
            # 根据输入的形状初始化 Linear 层
            self.linear = nn.Linear(in_features=data.shape[-1], out_features=self.out_features)
        return self.linear(data)


if __name__ == '__main__':
    model = LazyLinear()
    print(f'初始化前：\n{model}')
    
    # 第一次前向传播，触发参数初始化
    output = model(torch.randn(5, 20))
    print(f'初始化并前向传播：\n输出形状 = {output.shape}')
    print(f'初始化后：\n{model}')

输出结果：

1
2
3
4
5
6
7
8
9
10
初始化前：
LazyLinear()

初始化并前向传播：
输出形状 = torch.Size([5, 10])

初始化后：
LazyLinear(
  (linear): Linear(in_features=20, out_features=10, bias=True)
)

改进版实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import torch
import torch.nn as nn


class ImprovedLazyLinear(nn.Module):
    def __init__(self, out_features):
        super().__init__()
        self.out_features = out_features
        self.linear = None
        self._initialized = False

    def _initialize(self, input_shape):
        """根据输入形状初始化参数"""
        if not self._initialized:
            in_features = input_shape[-1]
            self.linear = nn.Linear(in_features, self.out_features)
            self._initialized = True
            
            # 将新创建的参数注册到模块中
            self.add_module('linear', self.linear)

    def forward(self, x):
        if not self._initialized:
            self._initialize(x.shape)
        return self.linear(x)

    def extra_repr(self):
        return f'out_features={self.out_features}, initialized={self._initialized}'

3.2 PyTorch 官方实现

LazyModuleMixin 类

PyTorch 从 1.8 版本开始，向torch.nn模块中引入了LazyModuleMixin类：

LazyModuleMixin 特性：

  • nn.Module中的层通过继承LazyModuleMixin类获得延后初始化特性
  • 在接收到数据并前向传播时，自动推断in_features参数并初始化
  • 这些有延后初始化特性的类以Lazy...为前缀命名

使用官方 LazyLinear

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import torch
import torch.nn as nn


class LazyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.LazyLinear(out_features=64)  # 使用延后初始化线性层
        self.relu = nn.ReLU()
        self.linear2 = nn.LazyLinear(out_features=10)

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x


if __name__ == '__main__':
    # 可以取消注释下面这行来抑制警告消息
    # import warnings; warnings.filterwarnings("ignore", category=UserWarning)

    model = LazyMLP()
    print(f'初始化前：\n{model}')
    
    # 第一次前向传播
    output = model(torch.randn(5, 20))
    print(f'初始化并前向传播：\n输出形状 = {output.shape}')
    print(f'初始化后：\n{model}')

输出结果：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
初始化前：
LazyMLP(
  (linear1): LazyLinear(in_features=0, out_features=64, bias=True)
  (relu): ReLU()
  (linear2): LazyLinear(in_features=0, out_features=10, bias=True)
)

初始化并前向传播：
输出形状 = torch.Size([5, 10])

初始化后：
LazyMLP(
  (linear1): Linear(in_features=20, out_features=64, bias=True)
  (relu): ReLU()
  (linear2): Linear(in_features=64, out_features=10, bias=True)
)

版本兼容性说明

开发状态警告：

截至 PyTorch 的 2.3.1 版本，延后初始化相关的模块仍处于活跃的开发状态，API 和功能可能随时发生修改。

在生产环境中使用时需要注意版本兼容性问题。可以通过设置警告过滤器来抑制相关警告消息。

4. 延后初始化的高级应用

4.1 复杂网络结构中的应用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import torch
import torch.nn as nn


class AdaptiveNetwork(nn.Module):
    def __init__(self, hidden_sizes, num_classes):
        super().__init__()
        self.hidden_sizes = hidden_sizes
        self.num_classes = num_classes
        
        # 构建延后初始化的层序列
        layers = []
        
        # 第一个隐藏层
        layers.append(nn.LazyLinear(hidden_sizes[0]))
        layers.append(nn.ReLU())
        
        # 中间隐藏层
        for i in range(1, len(hidden_sizes)):
            layers.append(nn.LazyLinear(hidden_sizes[i]))
            layers.append(nn.ReLU())
        
        # 输出层
        layers.append(nn.LazyLinear(num_classes))
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.network(x)


# 使用示例
model = AdaptiveNetwork(hidden_sizes=[128, 64, 32], num_classes=10)
print("模型定义完成，参数尚未初始化")

# 使用不同维度的输入测试
inputs = [
    torch.randn(10, 50),   # 50维输入
    torch.randn(10, 100),  # 100维输入
    torch.randn(10, 200),  # 200维输入
]

for i, input_tensor in enumerate(inputs):
    # 每次使用都需要重新创建模型，因为参数已经固定
    if i > 0:
        model = AdaptiveNetwork(hidden_sizes=[128, 64, 32], num_classes=10)
    
    output = model(input_tensor)
    print(f"输入维度: {input_tensor.shape[-1]}, 输出形状: {output.shape}")

4.2 与其他技术的结合

结合批归一化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class LazyBatchNormLinear(nn.Module):
    def __init__(self, out_features, use_bn=True):
        super().__init__()
        self.out_features = out_features
        self.use_bn = use_bn
        
        self.linear = nn.LazyLinear(out_features)
        if use_bn:
            self.bn = nn.LazyBatchNorm1d()
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.linear(x)
        if self.use_bn:
            x = self.bn(x)
        x = self.relu(x)
        return x

结合 Dropout

1
2
3
4
5
6
7
8
9
10
11
12
class LazyDropoutLinear(nn.Module):
    def __init__(self, out_features, dropout_rate=0.5):
        super().__init__()
        self.linear = nn.LazyLinear(out_features)
        self.dropout = nn.Dropout(dropout_rate)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        x = self.dropout(x)
        return x

5. 最佳实践与注意事项

5.1 使用建议

开发阶段建议：

  • 在开发和调试阶段，先用固定维度测试模型逻辑
  • 确保模型结构正确后再引入延后初始化
  • 使用多种不同维度的输入进行测试
  • 注意检查参数初始化后的模型状态

生产环境建议：

  • 在生产环境中谨慎使用，确保版本兼容性
  • 建议在模型部署前进行充分的测试
  • 考虑使用传统初始化方式以获得更好的稳定性
  • 监控模型的内存使用情况

调试技巧：

  • 使用model.named_parameters()检查参数初始化状态
  • 在关键位置添加断点检查张量形状
  • 使用torch.jit.trace验证模型的计算图
  • 记录不同输入维度下的性能表现

5.2 性能考虑

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import time
import torch
import torch.nn as nn


def benchmark_initialization():
    """比较传统初始化和延后初始化的性能"""
    
    # 传统初始化
    start_time = time.time()
    traditional_model = nn.Sequential(
        nn.Linear(1000, 512),
        nn.ReLU(),
        nn.Linear(512, 256),
        nn.ReLU(),
        nn.Linear(256, 10)
    )
    traditional_init_time = time.time() - start_time
    
    # 延后初始化
    start_time = time.time()
    lazy_model = nn.Sequential(
        nn.LazyLinear(512),
        nn.ReLU(),
        nn.LazyLinear(256),
        nn.ReLU(),
        nn.LazyLinear(10)
    )
    lazy_init_time = time.time() - start_time
    
    print(f"传统初始化时间: {traditional_init_time:.6f}s")
    print(f"延后初始化时间: {lazy_init_time:.6f}s")
    
    # 测试首次前向传播时间
    input_data = torch.randn(32, 1000)
    
    start_time = time.time()
    _ = traditional_model(input_data)
    traditional_forward_time = time.time() - start_time
    
    start_time = time.time()
    _ = lazy_model(input_data)  # 这里会触发参数初始化
    lazy_forward_time = time.time() - start_time
    
    print(f"传统模型首次前向传播时间: {traditional_forward_time:.6f}s")
    print(f"延后初始化模型首次前向传播时间: {lazy_forward_time:.6f}s")


if __name__ == '__main__':
    benchmark_initialization()

5.3 常见的延后初始化层

PyTorch 中的延后初始化层：

  • nn.LazyLinear：延后初始化的线性层
  • nn.LazyConv1d、nn.LazyConv2d、nn.LazyConv3d：延后初始化的卷积层
  • nn.LazyBatchNorm1d、nn.LazyBatchNorm2d、nn.LazyBatchNorm3d：延后初始化的批归一化层
  • nn.LazyInstanceNorm1d、nn.LazyInstanceNorm2d、nn.LazyInstanceNorm3d：延后初始化的实例归一化层

这些层都会在第一次前向传播时根据输入自动推断并初始化相应的参数。

总结

延后初始化是深度学习中的一项重要技术：

  1. 核心优势：动态推断参数维度，提高模型灵活性
  2. 应用场景：输入维度不确定、大型模型内存优化
  3. 实现方式：手动实现或使用PyTorch官方LazyModule
  4. 注意事项：调试难度增加，版本兼容性需要关注

延后初始化为构建灵活、高效的深度学习模型提供了新的思路，但在使用时需要权衡其带来的便利性和潜在的复杂性。在选择是否使用延后初始化时，应该根据具体的应用场景和需求来决定。

[1]
Dive into Deep Learning (D2L)