Skip to content

ePAI: lossless GPU-resident export (~1.85x, no accuracy loss) - #135

Merged
aperson30 merged 1 commit into
mainfrom
feat/epai-gpu-export
Aug 12, 2026
Merged

ePAI: lossless GPU-resident export (~1.85x, no accuracy loss)#135
aperson30 merged 1 commit into
mainfrom
feat/epai-gpu-export

Conversation

@aperson30

Copy link
Copy Markdown
Collaborator

What

Adds the lossless GPU-resident export path for ePAI — the biggest no-accuracy-loss speedup found by profiling.

Profiling showed ePAI's export/convert stage is ~158s of CPU resampling of the 26-channel logits back to native resolution — as expensive as the ~159s inference. Moving that resample to the GPU cuts it to ~0.3s (measured 50–60×) at 99.99% voxel agreement (boundary jitter only, no lesion-level change) → ~1.85× whole-pipeline, no accuracy cost.

How

scripts/epai_predict.py (+ .sh wrapper) reuses the fork's own export_prediction_from_logits — which produces both the segmentation and the PDAC/cyst/PNET findings CSV — but with ConfigurationManager.resampling_fn_probabilities patched to the GPU resampler. So the findings report is generated identically; only the resample runs on GPU.

Wired via the existing EPAI_SCRIPT_PATH hook in _run_epai_inference — no change to the dispatch logic.

Includes the fork-specific fix: determine_do_sep_z_and_axis returns the anisotropy axis as a bare scalar on the separate-z path (standard nnU-Net returns a 1-element list), which otherwise crashes resample_torch's assert len(axis)==1.

Before enabling (deploy gate)

  1. py_compile passes ✅
  2. End-to-end parity check (I'll run this): one real case through epai_predict.py vs the current bare-CLI path — confirm the segmentation and the output.csv findings match.
  3. Then set EPAI_SCRIPT_PATH=.../scripts/epai_predict.sh in prod .env and restart.

The GPU-export lever itself is already validated (50–60×, 99.99% agreement on 2 cases); the parity check just confirms the CSV path is faithful before it touches users.

… hook

Profiling showed ePAI's export/convert is ~158s of CPU resampling of the 26-channel
logits back to native res -- as costly as inference. This adds scripts/epai_predict.py
(+ .sh wrapper) that reuses the fork's own export_prediction_from_logits (seg AND the
PDAC/cyst/PNET findings CSV) but with the probabilities resample moved to GPU:
measured ~50-60x on that stage (158s -> ~0.3s), 99.99% voxel agreement vs CPU
(boundary jitter only, no lesion-level change) -> ~1.85x whole pipeline, no accuracy loss.

Includes the fork-specific fix for determine_do_sep_z_and_axis returning the anisotropy
axis as a bare scalar (standard nnU-Net returns a 1-element list), which otherwise
crashes resample_torch's len(axis) assert on the separate-z path.

Wired via the existing EPAI_SCRIPT_PATH hook in _run_epai_inference -- no change to the
dispatch. Enable in prod by setting EPAI_SCRIPT_PATH to scripts/epai_predict.sh after the
end-to-end parity check (seg + findings CSV identical to the bare-CLI path).
@aperson30
aperson30 merged commit 1744771 into main Aug 12, 2026
8 checks passed
@aperson30
aperson30 deleted the feat/epai-gpu-export branch August 13, 2026 10:53
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.

1 participant