ecmwf / anemoi-training

Apache License 2.0
16 stars 18 forks source link

Callbacks recompute predictions #117

Open HCookie opened 1 week ago

HCookie commented 1 week ago

What happened?

When using the rollout callbacks together, a number of predictions are done multiple times, wasting compute. Both RolloutEval and LongRolloutPlots call the rollout step independently.

What are the steps to reproduce the bug?

Run model with both RolloutEval and LongRolloutPlots enabled

Version

develop, v0.2.1+

Platform (OS and architecture)

ATOS

Relevant log output

No response

Accompanying data

No response

Organisation

ECMWF

Rilwan-Adewoyin commented 1 week ago

?? I propose having a parent HolderCallback class tailored to each type of task:

Each of these classes would contain task-specific versions of the standard callbacks. For instance, RolloutEval would include PlotLossForecast, while ReconstructionEval would use PlotLossReconstruction. These PlotLoss callbacks could inherit from a general PlotLoss class.

In each parent HolderCallback class, we would manage the forward loop and pass outputs to the relevant task-specific callbacks.