Refactor trainers

[adamjstewart]

In #1195 we introduced a new structure for our trainers. This PR updates all existing trainers to match.

• Add new base class (reduces code duplication)
• No *args or **kwargs (prevents argument typos, adds default values and type hints, required for LightningCLI)
• No typing.cast (declare type when initialized)
• __init__ first (should be first thing in docs)
• Added configure_(losses|metrics|models) methods (easier to override in subclass)

This is a great time to bring #996 back up. At this point, our "trainers" are referred to as:

• tasks: they all have a Task suffix
• trainers: they live in torchgeo.trainers
• models: this is what Lightning refers to them as
• modules: our YAML config file list them as module: (this will go away in my next PR)

I would love it if we could decide on a single naming scheme and be consistent...

Closes #1393

[adamjstewart]

BYOLTask has never had a backbone attribute. This line of code didn't do anything, and the tests are still slower and require more memory than necessary.

[adamjstewart]

Previously this would use a pretrained backbone and download weights from the internet by default, which doesn't match the behavior of any other trainer.

[adamjstewart]

Verbose isn't an option

[adamjstewart]

Many of these parameters were not previously documented.

[adamjstewart]

None of our other trainers require special handling for metric logging. Is it only because of Lightning-AI/torchmetrics#1832 (comment)?

[nilsleh]

• Should we have a public configure_models, configure_losses, configure_metrics for every trainer to make it easier for people to override these? This would be similar to configure_optimizers.

Regarding this point, for another domain library project I am also writing "trainers" that need a model, loss and metrics. When Isaac told me about hydra configs and instantiation I started using that at least for the model, loss and optimizer part. So you would have something like:

class BaseModel(LightningModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: type[torch.optim.Optimizer],
        loss_fn: nn.Module,
    ) -> None:
        self.model = model
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.loss_fn = loss_fn

    def configure_optimizers(self) -> dict[str, Any]:
        """Initialize the optimizer."""
        optimizer = self.optimizer(params=self.parameters())
        return optimizer

You could also provide defaults. But what I like is that you have some nice control via config files, so I could have:

method:
  _target_: BaseModel
  model:
    _target_: some.MLP # timm.create_model for timm models
    # all arguments to initiate the mlp for example
    n_outputs: 1
    n_hidden: [50]
  optimizer:
    _target_: torch.optim.Adam # can change optimizers here easily
    _partial_: true
    lr: 0.003
  loss_fn:
    _target_: torch.nn.MSELoss

So just via config you have control about optimizer and model architecture. Not sure if this is good practice but it has been convenient for running experiments and keeping track of configurations.

[adamjstewart]

You're in luck. Once this PR is merged I'll open a follow-up PR that switches everything to LightningCLI which supports exactly what you're describing without even changing the code:

• https://lightning.ai/docs/pytorch/stable/cli/lightning_cli_intermediate_2.html#multiple-optimizers

This all relies on jsonargparse which supports command-line, YAML, and JSON configuration.