Refactor and generalize loss.py

[ppegolo]

Fixes #629

Refactoring of the loss modules as discussed in #629 (with @SanggyuChong).

• dynamic lookup and registration as suggested by @HaoZeke
• moved <target_name> as top level section in the loss field in the input file
• grouped "torch-like" losses such as MSELoss, L1Loss, HuberLoss to a common interface

Contributor (creator of pull-request) checklist

• Tests updated (for new features and bugfixes)?
• Documentation updated (for new features)?
• Issue referenced (for PRs that solve an issue)?

Reviewer checklist

• CHANGELOG updated with public API or any other important changes?

---

📚 Documentation preview 📚: https://metatrain--635.org.readthedocs.build/en/635/

[HaoZeke]

Thanks for the PR @ppegolo! I think it looks great, have some comments for discussion inline, one thing which is missing for approval though is tests

[HaoZeke]

This could go to the constructor (__init__) so it isn't called every time compute is used..

[ppegolo]

It's called there because we have access to the targets only at compute, not at the time the loss is initialized in the trainer. Maybe there's a way to redesign it, but we've implemented it this way to have minimal change with respect to the current implementation of Trainer

[HaoZeke]

If we're always composing / delegatin then it might be nicer to store the object and use it, e.g.
self._base_loss = base_loss

[HaoZeke]

Just in case we forget to delegate anything else later.

[HaoZeke]

I think perhaps, if we're only discovering internal plugins, then an enum might be clearer (single source of truth). So something like:

class CoreLoss(Enum):
  MSE = TensorMapMSELoss
...

If / when we want to allow third party loss functions (i.e. in a separate pip install-able package) then we'd add in a Registry then which discovers plugins off of entry_points(group='metatrain.losses') and provides a unified view of all the loss functions present?

[HaoZeke]

Note that this is a minor design nit, I think we can easily do that later too... mostly I'm thinking an explicit interface for internal losses is clearer (and in keeping with the "Zen of Python") rather than auto-discovery here.

[Luthaf]

enum or even just a (factory) function somewhere doing

def get_loss(name, hypers):
    if name == "mse":
        return TensorMapMSELoss(**hypers)
    ...

We can then extend the factor however we want down the line

[ppegolo]

Thanks for the comments! If we need to have a register at some point when we'll allow for external loss functions, is it so bad to have it already now, to avoid having to re-implement it then?

[Luthaf]

I'm not sure we'll want to use a register for external loss function. The goal is not so much to have user A use the loss function provided by user B in some other python package, but rather allow user A to directly provide a python script with a loss. This can be done without any kind of global registry

[HaoZeke]

FWIW regarding external package entrypoint registration there's also this (draft) spec. https://scientific-python.org/specs/spec-0002/

[HaoZeke]

Suggested change

	# Use explicit registry_name if given, else snake_case from class name
	key = getattr(cls, "registry_name", None)
	if key is None:
	key = "".join(
	f"_{c.lower()}" if c.isupper() else c for c in name
	).lstrip("_")
	# only register the very first class under each key
	mcs._registry.setdefault(key, cls)
	if name != "LossInterface" and issubclass(cls, LossInterface):
	if "registry_name" not in attrs:
	raise TypeError(f"Class '{name}' must define a 'registry_name' class attribute")

Indentation might be off. It is good to have a fallback in general, but not if we only have a set number of loss functions in metatrain, since we can easily enforce a fixed name for each new loss (esp. with the enum suggestion).

If we end up with third party loss support, then this would make much more sense :)

[jwa7]

@Luthaf what's your opinion on this?

[Luthaf]

I agree that names for the losses should be explicit, but I don't think we need to go through through a registration mechanism. For now, doing something like

def get_loss(name: str):
    if name == "mse":
        return MSELoss(...)
    elif name == "mae":
        return MAELoss(...)
    elif name == "whatever":
        return WhateverLoss(...)
    else:
        raise ValueError(f"unknown loss function {name}")

The main advantage of a registry would be to allow loss defined outside of metatrain, but this is not something we need yet, so let's apply YAGNI

[PicoCentauri]

This has to go in any of the models that use the new loss function, right. Maybe it makes sense that we write a general json schema for the loss and only "import" those in the architecture ones.

I will figure out if there is a clever way to do it.