Auto updating custom layer parameters while backpropagating in pytorch - neural-network

Question

I have a pytorch custom layer defined as:
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
I then define my network where the custom layer is used:
class MyNet(nn.Module):
def __init__(self):
super(MyNet, self).__init__()
self.layer1 = MyCustomLayer()
def forward(self, x):
x = self.layer1(x)
return x
Now Let's instantiate MyNet and observe the issue:
# 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()
Iterating through model parameters shows None for custom layer parameter:
for p in model.parameters():
print (p.grad)
None
while directly accessing that parameter shows the correct grad value:
print(model.layer1.my_parameter.grad)
tensor([-1.4370])
This, in turn, prevents the optim step from updating the inner parameters automatically and leaves me with the hassle of having to update those manually. Anyone knows how I can address this issue?

Answer 1

What you did i.e. return x*self.my_registered_parameter[0] worked because you use the registered param for calculating the gradient.
When you call nn.Parameter it returns a new object and hence self.my_parameter that you use for the operation and the one registered are not same.
You can fix this by declaring the my_parameter as nn.Parameter
self.my_parameter = nn.Parameter(torch.rand(1, requires_grad = True))
self.my_registered_parameter= nn.ParameterList([self.some_parameter])
or you don't need to create my_registered_parameter variable at all. When you declare self.my_parameter as nn.Parameter it gets registered as a parameter.

Answer 2

Alright!
I had to switch the parameter variable calls within the custom layer to the nn.ParameterList object (i.e. return x*self.my_registered_parameter[0] instead of x*self.my_parameter ). In this example that meant changing the custom layer's parameter call in forward method to:
def forward(self, x):
return x*self.my_registered_parameter[0]
This is where it would've been nice to have pass by reference!
Now optim updates all the parameters as expected!