inference and validation performance fixes by maxwest-uw · Pull Request #1077 · lincc-frameworks/hyrax

maxwest-uw · 2026-10-14T22:46:03Z

Change Description

Some performance boosts for validation and inference.

This is a subset of a longer list of problems that claude identified as issues (it also discovered a problem and fix that was on @drewoldag 's performance investigation branch already and applied that as well, listed as 0.)

Claude identified problems:

Inference runs without torch.no_grad(). create_evaluator / Infer.run never turn off gradients, unlike the tester, validator and infer_stream. It makes no measurable difference on this tiny CNN, but costs real time and memory on realistic models.
The data cache fills up during inference even though nothing ever reads from it. It's a single pass, so it only costs time (including dead _data_size bookkeeping) and keeps the whole inference dataset in memory. At 10M+ objects that becomes a memory blowup.
Each data-loading setup re-parses the default config. _apply_configurations reads the default TOML again with the slow tomlkit parser every time a dataset is set up.
A per-batch loss.item() sync. This also means the reported validation loss, and the best-checkpoint score, is only the last batch's value, not the mean over the validation set.

  ┌───────────────────────────────────────────────────┬───────────────────────────────────┬───────────────┬───────────────┬───────────────────────────────────┐
  │                        Fix                        │          What I measured          │     main      │      Fix      │                 Δ                 │
  ├───────────────────────────────────────────────────┼───────────────────────────────────┼───────────────┼───────────────┼───────────────────────────────────┤
  │ #0 TensorBoard logging per item (new)             │ inference                         │ 2.674s        │ 1.474s        │ −44.9% (2.11× → 1.17× vs PyTorch) │
  ├───────────────────────────────────────────────────┼───────────────────────────────────┼───────────────┼───────────────┼───────────────────────────────────┤
  │                                                   │ validation, 0 / 2 workers         │ 3.76s / 3.83s │ 1.31s / 1.87s │ −65% / −51%                       │
  ├───────────────────────────────────────────────────┼───────────────────────────────────┼───────────────┼───────────────┼───────────────────────────────────┤
  │ #1 no_grad                                        │ inference                         │ 2.674s        │ 2.705s        │ no change (noise)                 │
  ├───────────────────────────────────────────────────┼───────────────────────────────────┼───────────────┼───────────────┼───────────────────────────────────┤
  │ #2 no cache during inference                      │ inference                         │ 2.674s        │ 2.272s        │ −15.0%                            │
  ├───────────────────────────────────────────────────┼───────────────────────────────────┼───────────────┼───────────────┼───────────────────────────────────┤
  │ #4 skip ConfigManager() re-parse                  │ inference                         │ 2.674s        │ 2.560s        │ −4.3%                             │
  ├───────────────────────────────────────────────────┼───────────────────────────────────┼───────────────┼───────────────┼───────────────────────────────────┤
  │ #9 average validation loss, one .item() per epoch │ validation, 0 workers             │ 3.755s        │ 3.417s        │ −9.0%                             │
  └───────────────────────────────────────────────────┴───────────────────────────────────┴───────────────┴───────────────┴───────────────────────────────────┘

  ┌───────────────────────┬────────┬────────┬───────────────────────────────────┐
  │                       │  main  │ Branch │                 Δ                 │
  ├───────────────────────┼────────┼────────┼───────────────────────────────────┤
  │ hyrax infer           │ 2.674s │ 1.326s │ −50.4% (2.11× → 1.04× vs PyTorch) │
  ├───────────────────────┼────────┼────────┼───────────────────────────────────┤
  │ validation, 0 workers │ 3.755s │ 1.286s │ −65.8%                            │
  └───────────────────────┴────────┴────────┴───────────────────────────────────┘

Solution Description

Code Quality

I have read the Contribution Guide and agree to the Code of Conduct
My code follows the code style of this project
My code builds (or compiles) cleanly without any errors or warnings
My code contains relevant comments and necessary documentation

DataProvider.resolve_data logged a duration for every item (including cache hits) and DataCache._do_insert logged one for every insert. Each is a SummaryWriter.add_scalar call, which made data loading the dominant cost of inference and validation (CIFAR-10 benchmark vs main: hyrax infer 2.67s -> 1.47s; validation over 3 epochs 3.76s -> 1.31s). The periodic cache count/size scalars (every 1000 inserts) are kept. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

create_evaluator set eval mode but never disabled autograd, so every infer_batch built a graph and kept activations for a backward pass that never happens. Wrap the evaluator's run in torch.no_grad(), matching create_tester, create_validator and InferStreamSession. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Infer and Test read every item exactly once, so caching only cost time (insert + size accounting) and held the whole dataset in memory. Add DataCache.disable() and call it on the providers those verbs create. (CIFAR-10 benchmark vs main: hyrax infer 2.67s -> 2.27s.) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

_apply_configurations constructed a ConfigManager just to call the static merge_configs. Construction re-reads and re-parses the default TOML with tomlkit, reads any hyrax_config.toml in the cwd, and runs migrations, once per DataProvider with a dataset_config. (Benchmark: infer 2.67s -> 2.56s.) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

validate_batch returned loss.item(), forcing a GPU->CPU sync on every batch, and only the last batch's value survived in validator.state.output, so the reported validation loss and best-checkpoint score were last-batch values. create_validator now accumulates each metric over the epoch and converts to float once. Built-in models with a standalone validate_batch return loss.detach(); models returning floats keep working. (Benchmark, nw=0: validation 3.76s -> 3.42s.) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

codecov · 2026-10-14T22:54:47Z

Codecov Report

❌ Patch coverage is 87.50000% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.98%. Comparing base (240763c) to head (7d76299).

Files with missing lines	Patch %	Lines
src/hyrax/datasets/data_cache.py	85.71%	1 Missing ⚠️
src/hyrax/models/hsc_autoencoder.py	0.00%	1 Missing ⚠️
src/hyrax/models/hsc_dcae.py	0.00%	1 Missing ⚠️
src/hyrax/models/hyrax_autoencoder.py	0.00%	1 Missing ⚠️
src/hyrax/models/image_dcae.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1077      +/-   ##
==========================================
- Coverage   71.04%   70.98%   -0.06%     
==========================================
  Files          92       92              
  Lines        9349     9362      +13     
==========================================
+ Hits         6642     6646       +4     
- Misses       2707     2716       +9

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

drewoldag · 2026-10-15T20:28:52Z

            # Note that `merge_configs` makes a copy of self.config, so the original
            # config will not be modified.
-            return cm.merge_configs(base_config, tmp_config)
+            return ConfigManager.merge_configs(base_config, tmp_config)


drewoldag · 2026-10-15T20:32:12Z

-        if self._insert_count % self.logging_interval == 0 and self._insert_count != 0:
+        # Log cache growth periodically; logging on every insert is too slow for the data hot path.
+        if self._insert_count % self.logging_interval == 0:
+            prefix = self.__class__.__name__


Reasonable and quick fix - I also wonder if it would be worth the effort to make a resource metric reporter that is disabled by default. Since model metrics and system metrics are logged via the same logger, it requires disentangling them, and might require instantiating multiple loggers in verb. Could be way more trouble than it's worth.

drewoldag · 2026-10-15T20:32:41Z

-        if had_any_miss:
-            tensorboardx_logger.log_duration_ts(f"{prefix}/cache_miss_s", start_time)
-        else:
-            tensorboardx_logger.log_duration_ts(f"{prefix}/cache_hit_s", start_time)


yeah, this ☝️

drewoldag · 2026-10-15T20:35:49Z

        decoded = self.forward(data)
        loss = self.criterion(decoded, data)

-        return {"loss": loss.item()}


~~I'm unsure about this one. If it works and reduces run time - great.~~

I understand now, sorry about about the previous comment. This is a good find :)

drewoldag · 2026-10-15T20:36:49Z

+    @validator.on(Events.ITERATION_COMPLETED)
+    def accumulate_metrics(validator):
+        for name, value in validator.state.output.items():
+            value = value.detach() if isinstance(value, torch.Tensor) else value
+            metric_sums[name] = metric_sums.get(name, 0) + value
+        batch_count[0] += 1


Do something similar for training and checkpointing?

I tried this on the training engine and didn't see as much improvement, but there was improvement!

maxwest-uw and others added 6 commits October 14, 2026 15:08

Merge branch 'main' into msw/perf-inference-fixes

7d76299

drewoldag reviewed Oct 15, 2026

View reviewed changes

Navigation Menu

Repository navigation

inference and validation performance fixes - #1077

inference and validation performance fixes#1077
maxwest-uw wants to merge 6 commits into
mainfrom
msw/perf-inference-fixes

maxwest-uw commented Oct 14, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Oct 14, 2026

Uh oh!

drewoldag Oct 15, 2026

Uh oh!

drewoldag Oct 15, 2026

Uh oh!

drewoldag Oct 15, 2026

Uh oh!

drewoldag Oct 15, 2026 •

edited

Loading

Uh oh!

drewoldag Oct 15, 2026

Uh oh!

drewoldag Oct 16, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

maxwest-uw commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Change Description

Solution Description

Code Quality

Uh oh!

codecov Bot commented Oct 14, 2026

Codecov Report

Uh oh!

drewoldag Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

drewoldag Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

drewoldag Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

maxwest-uw commented Oct 14, 2026 •

edited

Loading

drewoldag Oct 15, 2026 •

edited

Loading