Hello,
I wish to create intermediate trainable parameters in my network. I am able to collect intermediate gradients with the help of the retain_grad() method, and perform updates manually, but I’d prefer to work with the optim module PyTorch provides so that I don’t have to write my own functions for the various gradient updates there are out there.
I’m relatively new to PyTorch, and my understanding is that the updates on the intermediate params are not working since they are non leaf nodes. Is there a way to detach previous computations to create these leaf nodes, so that optimizer.step() can perform updates successfully? Or is there a way to pass these intermediate gradients into the optimizer object itself?
The code below gives a vague illustration of what I am trying to accomplish. Any help is greatly appreciated!
network = MyNetwork()

fc_weights = fc_weights = torch.tensor([[1,0],[0,1]],requires_grad=False,dtype=torch.float)
input = torch.tensor([[2],[3]],requires_grad=False,dtype=torch.float) # sample output of several layers that are frozen

params = torch.tensor([[0.25],[0.125]],requires_grad=True,dtype=torch.float) # learnable parameters for first final layer
params2 = torch.rand(2,1,requires_grad=True) # learnable parameters for second final layer
params3 = torch.rand(2,1,requires_grad=True) # learnable parameters for third final (final) layer

learning_rate = 1000 # just some sample learning rate to probe changes
optimizer = optim.SGD([params], lr=learning_rate)
optimizer2 = optim.SGD([params2],lr=learning_rate)
optimizer3 = optim.SGD([params3],lr = learning_rate)
optimizer.zero_grad() # zero gradients initially to be safe
optimizer2.zero_grad()
optimizer3.zero_grad()

print(params) # check value before nudge
print(params2) # check value before nudge
print(params3) # check value before nudge
print()

x = torch.hstack((input,params))
params2 = network(x,fc_weights)
params2.retain_grad() # keep intermediate grads
print(params2.is_leaf) # prints False
params3 = network(params2,fc_weights)
params3.retain_grad() # keep intermdiate grads
print(params3.is_leaf) # prints False
print()

loss = MyLoss(params3)
loss.backward()

print(params.grad) # intermediate gradients successfully computed
print(params2.grad) # intermediate gradients successfully computed
print(params3.grad) # intermediate gradients successfully computed
print()

optimizer.step() # gradient update/ nudge
optimizer2.step() # gradient update/ nudge
optimizer3.step() # gradient update/ nudge

print(params) # this learnable parameter is updated successfully because it is a leaf node.
print(params2) # this learnable parameter is not updated, which is what I am trying to figure out
print(params3) # this learnable parameter is not updated, which is what I am trying to figure out