如何用我自己的函数初始化pytorch中的权重？（cs231n PyTorch 笔记本）

Question

在cs231n pytorch笔记本中，他们定义了random_weight函数来返回通过Kaiming归一化初始化的张量对象，以及zero_weight函数来返回包含零偏差的张量。

random_weight 实现：

def random_weight(shape):
    if len(shape) == 2:  # FC weight
        fan_in = shape[0]
    else:
        fan_in = np.prod(shape[1:]) 
    w = torch.randn(shape, device=device, dtype=dtype) * np.sqrt(2. / fan_in)
    w.requires_grad = True
    return w

零权重实现：

def zero_weight(shape):
    return torch.zeros(shape, device=device, dtype=dtype, requires_grad=True)

现在的目标是使用 pytorch 顺序 api 定义一个三层 ConvNet，我们必须使用上述两个函数来初始化 ConvNet 中的所有权重。

我发现了一些像这样的实现：

model = nn.Sequential(....)

def init_weights(layer):
    if type(layer) in [nn.Conv2d, nn.Linear]:
        random_weight(layer.weight.size())
        zero_weight(layer.bias.size())

model.apply(init_weights)
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, nesterov=True)

但我怀疑权重不是由这两个函数初始化的。因为我在调用random_weight之前和之后打印了layer.weight.data，所以我得到了相同的数据。

所以我尝试像这样修改 init_weights 代码：

def init_weights(layer):
    if type(layer) in [nn.Conv2d, nn.Linear]:
        layer.weight.data = random_weight(layer.weight.size())
        layer.bias.data = zero_weight(layer.bias.size())

但性能确实很糟糕，我在验证集上得到了大约 10% 的准确率。那么，有谁知道使用这两个函数来初始化权重的正确方法？

欢迎所有建议！多谢！