Register autograd - missing forward gradients

[bananenpampe]

Hello,

I am receiving this error when using the register autograd function:

ValueError: expected the precomputed TensorMap to contain gradients with respect to 'positions' since one of the system ``requires_grad`` for its positions

I was hoping that the register autograd function could actually be used to register precomputed features using rascaline.torch after having called .backward(retain_graph=False) which detaches feats from the computational graph that system.positions are part of.

This is the minimal example:

import ase
import ase.build
import rascaline.torch

frame = ase.build.bulk("Al", "fcc", a=4.05, cubic=True)
system = rascaline.torch.systems_to_torch(frame,
                                 positions_requires_grad=True,
                                 cell_requires_grad=False)

HYPERS = {
    "cutoff": 5.,
    "max_radial": 2,
    "max_angular": 1,
    "atomic_gaussian_width": 0.3,
    "center_atom_weight": 0.0,
    "radial_basis": {
        "Gto": {},
    },
    "cutoff_function": {
        "ShiftedCosine": {"width":0.5},
    },
    "radial_scaling":{"Willatt2018": {"exponent": 6.0, "rate": 2.0, "scale": 2.0}}
}

calculator = rascaline.torch.calculators.SoapPowerSpectrum(**HYPERS)

feats = calculator.compute(system)

rascaline.torch.register_autograd(system, feats)

[Luthaf]

You need to compute explicit gradients in calculator.compute(system), i.e. replace the call by calculator.compute(system, gradients=["positions"])

[bananenpampe]

Thank you! Does this result in an increased memory consumption of the feats, after potential densification?

[Luthaf]

It should not, especially if you remove the gradient blocks with

rascaline.torch.register_autograd(system, feats)
feats = metatensor.remove_gradients(feats)

If you don't remove the gradients, it might increase memory.

[bananenpampe]

okay, but will it be possible to train/evaluate a model using double backprop (I am thinking about forces here) using these lines ?:

rascaline.torch.register_autograd(system, feats)
feats = metatensor.remove_gradients(feats)
loss(target, model(feats)).backward()

[Luthaf]

remove_gradients removes the explicit, forward mode gradient stored inside TensorBlock. detach is the function that would break the computational graph. So yes, this should be fine!

[bananenpampe]

okay, awesome, thank you very much

[bananenpampe]

closing, but I should add some documentation

[bananenpampe]

@E-Rum