ePAI: lossless GPU-resident export (~1.85x, no accuracy loss) - #135
Merged
Conversation
… 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).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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(+.shwrapper) reuses the fork's ownexport_prediction_from_logits— which produces both the segmentation and the PDAC/cyst/PNET findings CSV — but withConfigurationManager.resampling_fn_probabilitiespatched to the GPU resampler. So the findings report is generated identically; only the resample runs on GPU.Wired via the existing
EPAI_SCRIPT_PATHhook in_run_epai_inference— no change to the dispatch logic.Includes the fork-specific fix:
determine_do_sep_z_and_axisreturns the anisotropy axis as a bare scalar on the separate-z path (standard nnU-Net returns a 1-element list), which otherwise crashesresample_torch'sassert len(axis)==1.Before enabling (deploy gate)
py_compilepasses ✅epai_predict.pyvs the current bare-CLI path — confirm the segmentation and theoutput.csvfindings match.EPAI_SCRIPT_PATH=.../scripts/epai_predict.shin prod.envand 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.