Define custom loss function

[tgmaxson]

Is there any mechanism to define a custom loss function component via python extension or something of the sort? Currently it seems they are hard coded to search nequip.nn and it unclear where exactly I would extend the system without direct modification of nequip.

To set some limits since I expect they would matter, lets say I want to follow some sort of ERF-like curve in loss function, where stresses start with a low weight and smoothly increase to some maximum. Ideally I would define this in some sort of simple math expression in the config.yaml. If this is not possible, can you suggest a minimum example of how this could be directly added to nequip?

[Linux-cpp-lisp]

Hi @tgmaxson ,

Yes, custom loss terms are possible, but a bit awkward in syntax at present.

Syntax is something like:

loss_coeffs:
      forces:
        - 0.1
        - !!python/object:my_extension_package.loss_terms.MyL1Loss {}

This funny syntax tells PyYAML to construct an object of that class and put it in the config.

The implementation, at simplest, looks like

class MyL1Loss:
    def __call__(
        self,
        pred: dict,
        ref: dict,
        key: str,
        mean: bool = True,
    ):
        assert mean
        pred = pred[key]
        ref = ref[key]
        term = (pred - ref).abs().sum()
        return term

key here is forces, pred and ref are the predicted and reference AtomicDataDict.Types, and mean can basically be ignored as long as you don't want to use the loss as a metric (if you do, return instead a per-(atom/graph/whatever) tensor of losses).

(See SimpleLoss in the repo for the actual corresponding implementation, which handles some more edge cases that you don't really need to think about.)