Skip to content

Average long_time_metrics time curves over validation batches#93

Open
Adonyth wants to merge 1 commit into
PolymathicAI:masterfrom
Adonyth:fix/long-time-metrics-batch-average
Open

Average long_time_metrics time curves over validation batches#93
Adonyth wants to merge 1 commit into
PolymathicAI:masterfrom
Adonyth:fix/long-time-metrics-batch-average

Conversation

@Adonyth

@Adonyth Adonyth commented Jul 12, 2026

Copy link
Copy Markdown

Fixes #78

Inside Trainer.validation_loop, the time curves collected for long_time_metrics are merged with time_logs |= new_time_logs, which overwrites every key on each batch — the curves handed to plot_all_time_metrics reflect only the last validation batch. A two-batch toy example makes it concrete: if a metric's curve is 100.0 on batch 1 and 1.0 on batch 2, the logged value is 1.0 instead of the 50.5 average. This is inconsistent with loss_dict three lines below, which correctly accumulates loss_dict.get(k, 0.0) + v / denom.

The fix mirrors the loss_dict accumulation exactly, so the time curves become means over the same denom (the number of validation batches actually processed):

if k in long_time_metrics or "spectral_error" in k:
    for log_name, log_value in new_time_logs.items():
        time_logs[log_name] = (
            time_logs.get(log_name, 0.0) + log_value / denom
        )

Test note: Trainer isn't constructible without a full model/datamodule stack, so tests/test_trainer.py builds one via Trainer.__new__, sets only the attributes validation_loop reads, stubs rollout_model to return pre-made (y_pred, y_ref) batches, and captures what reaches plot_all_time_metrics. It asserts the captured full_VRMSE_rollout curve equals the mean of the two per-batch curves (and not the last batch's curve). The test fails on current master and passes with this change; ruff check / ruff format are clean under the pre-commit pinned v0.6.4. Happy to restructure the test if you'd prefer a different fixture approach.

Found while running an independent quantified-impact study of the benchmark's metric design; happy to adjust to maintainer preferences.

Inside Trainer.validation_loop, 'time_logs |= new_time_logs' overwrote
the long_time_metrics time curves on every batch, so the curves passed
to plot_all_time_metrics reflected only the last validation batch.
Accumulate them as running means over the batch count instead, exactly
mirroring how loss_dict is accumulated a few lines below. Add a
lightweight validation_loop test with a stubbed rollout.

Fixes PolymathicAI#78
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: long_time_metrics reflect only last batch of data and not the average over all batches

1 participant