Hi!
I’m working on a project on molecule energies. The loss function consists of two parts: the energy term and the force term and forces are derivatives wrt the input coordinates. This is part of my code:
def energies_forces_trainer(container, optimizer, loss_fn):
def train_and_store_loss(engine, batch):
model.train()
optimizer.zero_grad()
inputs, targets = batch
final = {'energies': [], 'forces': []}
for s, c in inputs:
results = model((s, c)) # c stands for coordinates
_, energies = energy_shifter(results) # Some transformation
forces = -torch.autograd.grad(energies.sum(), c)[0] # I want to compute the result's derivatives wrt the input
final['energies'].append(results[1])
final['forces'].append(forces)
final['energies'] = torch.cat(final['energies'])
final['forces'] = torch.cat(final['forces'])
loss = loss_fn('energies')(targets, final) + loss_fn('forces')(targets, final) # Energy loss and force loss
loss.backward() # Error ocurrs
optimizer.step()
return loss.item()
return ignite.engine.Engine(train_and_store_loss)
This code produces ‘Trying to backward through the graph a second time’ error. So does torch.autograd.grad accumulate grads? How can I include loss of derivatives?
Thank you!