这个代码片段展示了如何用 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)
这些权重张量的生成方式如下:
python
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)
]
示例代码
下面是一个完整的示例,展示如何使用上述权重初始化方式初始化一个简单的神经网络:
python
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
方法中使用截断正态分布初始化权重。通过打印权重,我们可以验证它们是否按预期初始化。
说明
- 定义网络 :
CustomModel
包含四个线性层。第一层输入尺寸为x_dim
,输出尺寸为width + n_double
。接下来的两层也是同样的输出尺寸,最后一层输出尺寸为 1。 - 初始化权重 :在
init_weights
方法中,我们使用截断正态分布(通过torch.fmod
)初始化每一层的权重。我们对生成的正态分布取模 2,使得权重在 -2 和 2 之间。 - 打印参数 :我们通过
model.named_parameters()
方法遍历模型的参数,并打印每层参数的名称、尺寸和前两个值。
进一步说明
- 截断正态分布 :使用
torch.normal
生成正态分布的随机数,然后使用torch.fmod
将这些随机数的范围限制在 -2 到 2 之间。 - 超参数 :
x_dim
是输入的特征维度,width
是每层的宽度(即神经元数量),n_double
是一个附加参数,用于增加每层的输出维度。init_sd_first
、init_sd_middle
和init_sd_last
是每层权重初始化的标准差。
这个示例展示了如何使用截断正态分布初始化神经网络的权重,并打印每层的参数。如果您有更多问题或需要进一步的帮助,请告诉我!