这个代码片段展示了如何用 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_firstinit_sd_middleinit_sd_last 是每层权重初始化的标准差。

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

07-07 11:46