首页 > 解决方案 > 在 pytorch 中反向传播时自动更新自定义层参数

问题描述

我有一个 pytorch 自定义层定义为:

class MyCustomLayer(nn.Module):
  def __init__(self):
    super(MyCustomLayer, self).__init__()

    self.my_parameter = torch.rand(1, requires_grad = True)

    # the following allows the previously defined parameter to be recognized as a network parameter when instantiating the model
    self.my_registered_parameter = nn.ParameterList([nn.Parameter(self.my_parameter)])

  def forward(self, x):
    return x*self.my_parameter

然后我定义使用自定义层的网络:

class MyNet(nn.Module):
  def __init__(self):
    super(MyNet, self).__init__()
    self.layer1 = MyCustomLayer()

  def forward(self, x):
    x = self.layer1(x)
    return x

现在让我们实例化 MyNet 并观察问题:

# instantiate MyNet and run it over one input value
model = MyNet()
x = torch.tensor(torch.rand(1))
output = model(x)
criterion = nn.MSELoss()
loss = criterion(1, output)
loss.backward()

遍历模型参数显示None自定义层参数:

for p in model.parameters():
    print (p.grad)

None

直接访问该参数显示正确的grad值:

print(model.layer1.my_parameter.grad)

tensor([-1.4370])

这反过来又阻止了优化步骤自动更新内部参数,让我不得不手动更新这些参数。任何人都知道我该如何解决这个问题?

标签: neural-networkgradientpytorchbackpropagation

解决方案


好吧!我不得不将自定义层中的参数变量调用切换到nn.ParameterList对象(即return x*self.my_registered_parameter[0],而不是 x*self.my_parameter )。在此示例中,这意味着将自定义层在 forward 方法中的参数调用更改为:

  def forward(self, x):
    return x*self.my_registered_parameter[0]

这是通过引用传递的好地方!

现在 optim 按预期更新所有参数!


推荐阅读