Right now, I have a use case like following
A is a torch.tensor like following
A =
[[1,x,y],
[1,2,3],
[1,z,3]]
only elements x, y, z in A are differentiable, other elements are just constant.
For example, if the cost = tr(A.A)
cost = 14 + 2 x + 2 y + 6 z
When I do backtracking, I only want to differentiate and update with respect to x, y, z. Certainly this example is just a toy example, not the true complicated one.
How can I realize such use case?
The only way I see to achieve it is to use a mask tensor to zero-out the gradients of the elements you don't want to modify, like this :
A = torch.tensor([[1.0, 100, 100], [1, 2, 3], [1, 100, 3]])
# I have initialized your x,y,z to 100, so the mask should only have these set to 1
mask = (A==100).to(dtype=torch.float32)
#optimizer to backpropagate
opt = torch.optim.SGD([A], lr=1e-2)
# optimization loop
for i in range(10):
# backprop the loss gradient
(A.matmul(A)).diag().pow(2).sum().backward()
# mask the gradients
with torch.no_grad():
A.grad *= mask
# optimization step
opt.step()
A.grad.zero_()
print(A)
Only the x,y,z coordinates should be modified then