Hi, I’ve written a NN where the loss function consists of gradients of the output with respect to the inputs. The code works fine, but I want to inspect how the gradients flow when I’m updating the NN weights with respect to the loss function. I can’t figure out how the computation graph is laid out here. I hope someone can help me understand how exactly are the gradients computed.
Here’s my code
dtype = torch.float
device = torch.device("cuda:0") # Uncomment this to run on GPU
x = torch.tensor([[1., 1.]], device=device, dtype=dtype, requires_grad=True)
w1 = torch.tensor([[ 0.6561, 0.6202, -1.5620],
[ 1.1547, -1.3227, 0.4719]], device=device, dtype=dtype, requires_grad=True)
w2 = torch.tensor([[-0.4917],
[-0.3886],
[-1.4218]], device=device, dtype=dtype, requires_grad=True)
learning_rate = 0.01
z1 = x.mm(w1)
a1 = torch.sin(z1)
z2 = a1.mm(w2)
p = torch.sin(z2)
grads, = torch.autograd.grad(p, x,
grad_outputs=p.data.new(p.shape).fill_(1),
create_graph=True, only_inputs=True)
dpdx, dpdt = grads[:,0], grads[:,1]
pde = dpdt + dpdx
loss = pde.pow(2).mean()
loss.backward()
with torch.no_grad():
w1 -= learning_rate * w1.grad
w2 -= learning_rate * w2.grad
# Manually zero the gradients after updating weights
w1.grad.zero_()
w2.grad.zero_()
I just can’t understand how I get w1.grad and w2.grad (i.e. what are the connections in the computation graph) when I run loss.backward().