Hi @ptrblck. I have a similar problem where I’m creating a custom loss function extending the nn.module. However the results are absurd.
def forward(self, x):
with torch.set_grad_enabled(True):
time_step =torch.tensor(0.01)
out=self._rk4_step1(self.function, x, 0, time_step)
return out
def function(self,x,t):
self.n = n = x.shape[1]//2
qqd = x.requires_grad_(True)
L = self._lagrangian(qqd).sum()
J = grad(L, qqd, create_graph=True)[0] ;
DL_q, DL_qd = J[:,:n], J[:,n:]
DDL_qd = []
for i in range(n):
J_qd_i = DL_qd[:,i][:,None]
H_i = grad(J_qd_i.sum(), qqd, create_graph=True)[0][:,:,None]
DDL_qd.append(H_i)
DDL_qd = torch.cat(DDL_qd, 2)
DDL_qqd, DDL_qdqd = DDL_qd[:,:n,:], DDL_qd[:,n:,:]
T = torch.einsum('ijk, ij -> ik', DDL_qqd, qqd[:,n:])
qdd = torch.einsum('ijk, ij -> ik', DDL_qdqd.pinverse(), DL_q - T)
return torch.cat([qqd[:,self.n:], qdd], 1)
def _lagrangian(self, qqd):
x = F.softplus(self.fc1(qqd))
x = F.softplus(self.fc2(x))
# x = F.softplus(self.fc3(x))
L = self.fc_last(x)
return L
def _rk4_step1(self, f, x, t, h):
# one step of Runge-Kutta integration
k1 = torch.mul(f(x, t),h)
k2 = torch.mul(f(x + k1/2, t + h/2),h)
k3 = torch.mul(f(x + k2/2, t + h/2),h)
k4 = torch.mul(f(x + k3, t + h),h)
return x + 1/6 * (k1 + 2 * k2 + 2 * k3 + k4)
Is the autograd able to track all gradient even though I’m calling the forward multiple times in the rk4_step?