Hey, I am currently experiencing an issue with my regularization loss (which is calculated and back propagated through without first being combined with a normal loss) giving me the error ‘trying to backward through the graph a second…’.
I modify the weights using the outputs of another layer and that does not allow me to specify retain_graph=False in torch.autograd.grad().
Here is a simple code example which fails:
w = nn.Parameter(torch.zeros(10, 10))
dense = nn.Linear(10, 10)
x = torch.rand(1, 10).requires_grad_(True)
y = x.matmul(w.t() * dense(x).mean(0)).mean()
grad = torch.autograd.grad(
outputs=y,
inputs=x,
retain_graph=False,
create_graph=True,
only_inputs=True
)[0]
grad.mean().backward()
And here is a version that works but without modifying the weight.
w = nn.Parameter(torch.zeros(10, 10))
x = torch.rand(1, 10).requires_grad_(True)
y = x.matmul(w.t()).mean()
grad = torch.autograd.grad(
outputs=y,
inputs=x,
retain_graph=False,
create_graph=True,
only_inputs=True
)[0]
grad.mean().backward()
Is there any way for me to run this without using retain_graph=True? Or is there no difference in memory/performance from retaining the graph for my specific example?