nn.modules.lazy.LazyModuleMixin用法详解

torch.nn.modules.lazy.LazyModuleMixin 是 PyTorch 中的一个混合类（mixin），它用于创建那些延迟初始化参数的模块，也就是“懒加载模块（lazy modules）”。这些模块从它们的第一次前向传播输入中推导出参数的形状。在第一次前向传播之前，它们包含 torch.nn.UninitializedParameter，这些参数不应被访问或使用；在之后，它们包含常规的 torch.nn.Parameter。

用途

懒加载模块的主要用途是简化网络的构建过程，因为它们不需要计算一些模块参数，比如 torch.nn.Linear 中的 in_features 参数。这对于处理具有可变输入大小的数据非常有用。

参数

• *args 和 **kwargs：用于初始化混合类的任何标准参数。

使用技巧和注意事项

1. 转换数据类型和设备：在构建含有懒加载模块的网络后，首先应该将网络转换为所需的数据类型（dtype）并放置在预期的设备上。这是因为懒加载模块只执行形状推断，因此常规的数据类型和设备放置行为适用。
2. 执行“干运行”：在使用网络之前，应该执行“干运行”以初始化模块中的所有组件。这些“干运行”通过网络发送正确大小、数据类型和设备的输入，初始化每一个懒加载模块。
3. 初始化顺序变化：使用懒加载模块时，网络参数的初始化顺序可能会改变，因为懒加载模块总是在其他模块之后初始化。
4. 序列化和反序列化：懒加载模块可以像其他模块一样使用状态字典（state dict）进行序列化。但是请注意，如果在状态加载时参数已经初始化，进行“干运行”时不会替换这些参数。

示例代码

以下是一个使用 torch.nn.modules.lazy.LazyModuleMixin 的示例：

import torch
import torch.nn as nn

class LazyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.LazyLinear(10)  # 懒加载线性层
        self.relu1 = nn.ReLU()
        self.fc2 = nn.LazyLinear(1)
        self.relu2 = nn.ReLU()

    def forward(self, input):
        x = self.relu1(self.fc1(input))
        y = self.relu2(self.fc2(x))
        return y

# 构建懒加载网络
lazy_mlp = LazyMLP()

# 转换网络的设备和数据类型
lazy_mlp = lazy_mlp.cuda().double()

# 执行干运行以初始化懒加载模块
lazy_mlp(torch.ones(10, 10).cuda())

# 在初始化后，LazyLinear模块变为常规Linear模块
print(lazy_mlp)

# 附加优化器
optim = torch.optim.SGD(lazy_mlp.parameters(), lr=0.01)

在这个例子中，LazyMLP 类中的两个 nn.LazyLinear 层在第一次前向传播时根据输入自动初始化。在干运行之后，这些层就变成了常规的 nn.Linear 层，可以用于后续的训练或推理。