【PyTorch】torch.fmod使用截断正态分布truncated normal distribution初始化神经网络的权重

这个代码片段展示了如何用 PyTorch 初始化神经网络的权重，具体使用的是截断正态分布（truncated normal distribution）。截断正态分布意味着生成的值会在一定范围内截断，以防止出现极端值。这里使用 torch.fmod 作为一种变通方法实现这一效果。

详细解释

1. 截断正态分布

截断正态分布是对正态分布的一种修改，确保生成的值在一定范围内。具体来说，torch.fmod 函数返回输入张量除以 2 的余数（即使得生成的值在 -2 到 2 之间）。

2. 权重初始化

代码中，四个权重张量按不同的标准差（init_sd_first, init_sd_middle, init_sd_last）从截断正态分布中生成。具体的维度分别是：

• 第一层的权重张量形状为 (x_dim, width + n_double)
• 中间层的两个权重张量形状为 (width, width + n_double)
• 最后一层的权重张量形状为 (width, 1)

这些权重张量的生成方式如下：

initial_weights = [
    torch.fmod(torch.normal(0, init_sd_first, size=(x_dim, width + n_double)), 2),
    torch.fmod(torch.normal(0, init_sd_middle, size=(width, width + n_double)), 2),
    torch.fmod(torch.normal(0, init_sd_middle, size=(width, width + n_double)), 2),
    torch.fmod(torch.normal(0, init_sd_last, size=(width, 1)), 2)
]

示例代码

下面是一个完整的示例，展示如何使用上述权重初始化方式初始化一个简单的神经网络：

import torch
import torch.nn as nn

class CustomModel(nn.Module):
    def __init__(self, x_dim, width, n_double, init_sd_first, init_sd_middle, init_sd_last):
        super(CustomModel, self).__init__()
        self.linear1 = nn.Linear(x_dim, width + n_double)
        self.linear2 = nn.Linear(width + n_double, width + n_double)
        self.linear3 = nn.Linear(width + n_double, width + n_double)
        self.linear4 = nn.Linear(width + n_double, 1)
        
        self.init_weights(init_sd_first, init_sd_middle, init_sd_last)

    def init_weights(self, init_sd_first, init_sd_middle, init_sd_last):
        self.linear1.weight.data = torch.fmod(torch.normal(0, init_sd_first, size=self.linear1.weight.size()), 2)
        self.linear2.weight.data = torch.fmod(torch.normal(0, init_sd_middle, size=self.linear2.weight.size()), 2)
        self.linear3.weight.data = torch.fmod(torch.normal(0, init_sd_middle, size=self.linear3.weight.size()), 2)
        self.linear4.weight.data = torch.fmod(torch.normal(0, init_sd_last, size=self.linear4.weight.size()), 2)

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

# 定义超参数
x_dim = 10
width = 20
n_double = 5
init_sd_first = 0.1
init_sd_middle = 0.1
init_sd_last = 0.1

# 初始化模型
model = CustomModel(x_dim, width, n_double, init_sd_first, init_sd_middle, init_sd_last)

# 打印权重以验证初始化
for name, param in model.named_parameters():
    if 'weight' in name:
        print(f"{name} initialized with values: \n{param.data}\n")

在这个示例中，我们定义了一个简单的神经网络 CustomModel，并在 init_weights 方法中使用截断正态分布初始化权重。通过打印权重，我们可以验证它们是否按预期初始化。

说明

1. 定义网络：CustomModel 包含四个线性层。第一层输入尺寸为 x_dim，输出尺寸为 width + n_double。接下来的两层也是同样的输出尺寸，最后一层输出尺寸为 1。
2. 初始化权重：在 init_weights 方法中，我们使用截断正态分布（通过 torch.fmod）初始化每一层的权重。我们对生成的正态分布取模 2，使得权重在 -2 和 2 之间。
3. 打印参数：我们通过 model.named_parameters() 方法遍历模型的参数，并打印每层参数的名称、尺寸和前两个值。

进一步说明

• 截断正态分布：使用 torch.normal 生成正态分布的随机数，然后使用 torch.fmod 将这些随机数的范围限制在 -2 到 2 之间。
• 超参数：x_dim 是输入的特征维度，width 是每层的宽度（即神经元数量），n_double 是一个附加参数，用于增加每层的输出维度。init_sd_first、init_sd_middle 和 init_sd_last 是每层权重初始化的标准差。

这个示例展示了如何使用截断正态分布初始化神经网络的权重，并打印每层的参数。如果您有更多问题或需要进一步的帮助，请告诉我！