Where the time and energy go in a kidney‑tumor segmentation model

A KiTS segmentation model takes an abdominal CT scan and returns a mask of each kidney and any tumor. It is accurate, but running one is not cheap: a single scan takes a few minutes of mixed CPU and GPU work, which at any real volume of patients adds up to meaningful energy and carbon. This is a walk through where that cost actually goes for one case, and then how far a sequence of inference‑side changes — none of which retrain the model — can bring it down. Unless noted, every figure is a mean per case, measured on an RTX 4090 over the 189 held‑out cases (00400–00588).

CPU work GPU work disk I/O + write

The baseline pipeline (B0) — load → CPU → GPU → CPU → export

A case moves CPU → GPU → CPU, and the colour of each box marks which device does the work. This is the stock configuration — call it B0 — with the mean time each step takes.

CT scan
.nii.gz
→
DISK · CPU

1 · Load

Read + gunzip the .nii.gz, parse the NIfTI header & voxels into RAM.

≈ 1–3 s
→
CPU

2 · Preprocess

Crop to non‑zero → resample to 0.78 mm isotropic (order‑3 spline, separate‑Z) → normalize (CT clip + z‑score). Single‑threaded scipy.

≈ 50 s
→
GPU

3 · Predict

Sliding window: tile into fixed patches → forward pass (fp32) → Gaussian‑weighted accumulate softmax → ×8 TTA mirror flips (8 passes per patch).

≈ 138 s
→
CPU / GPU

4 · Export

Softmax GPU→CPU → argmax to labels → resample back to original spacing → write .nii.gz (SimpleITK gzip).

≈ 35 s
→
mask
.nii.gz

End to end the baseline takes about 223 seconds per case, and the GPU prediction step is the largest single slice. The rest of this piece follows what happens to that budget — in both time and energy — as the configuration changes.

Where the time & energy go — the full ladder (RTX 4090 · n=189)

Each rung is a configuration change, noted on the left. The bars are the mean per‑case stage split, scaled to B0’s 223 s; at the right of each is that config’s mean seconds per case and, below it, its mean prediction energy (kJ per case, from the GPU‑locked NVML counter). The prediction segment is the cleanest signal to follow — it falls from 138 s to about 5 s down the ladder, and the energy falls with it. (CPU preprocess and export times vary from one pod to another, since each config ran on its own machine, so the prediction segment is the more reliable trend. All figures are means over the 189 cases.)

B0 stock: fp32 · 8‑TTA · step .50
50
predict 138s
35
223 s45.1 kJ
B0 + TRT + TensorRT forward only (lossless)
50
predict 67s
35
152 s17.3 kJ
O1 + fp16 · all‑in‑GPU (lossless)
70
predict 71s
59
199 s24.0 kJ
O3 + drop TTA · stride .75
32
25
62 s1.36 kJ
O4 + TensorRT engine
41
35
81 s0.92 kJ
O5 + restore 4‑pass TTA
50
11
40
101 s3.26 kJ
O6 + GPU‑resample · pipeline
pre 27
9
10
47 s3.07 kJ
O6 + parallel + parallel preprocess · measured L40S / 10‑case
pre
9
10
~18.6 s~39.7 kJ total · L40S
0 s112 s223 s
B0 + TRT: stock pipeline with only the forward pass swapped to a TensorRT engine — masks bit‑for‑bit identical (Dice 0.887, lossless). On a paired 10‑case 4090 run the forward ran 2.05× faster for 2.6× less predict energy (51.6 → 19.9 kJ/case); the rung above projects that onto the n=189 times. TRT helps most here because the 8 TTA passes make the forward pass dominate the work — with TTA off (O3→O4) the same engine saves only about 15 %.
O6’s GPU‑resample & pipelining: export has to resample the mask back to the scan’s native voxel grid. Stock nnU‑Net does this on the CPU with scipy — a single‑thread 3D cubic‑spline over the whole volume, tens of seconds. O6 moves it to the GPU with torch (parallel interpolation across thousands of cores), which brings it down to about a second, and overlaps the remaining gzip write with the next case (pipelining). On the output side this is lossless: argmax picks the winning label per voxel, so a tiny interpolation difference almost never flips the result. (Doing the same on the input side is not free — the resampled image feeds the whole network, and torch’s trilinear ≠ the scipy cubic‑spline nnU‑Net was trained on, so the masks change. That’s the separate “research” rung.)

The shift — the bottleneck moves off the GPU

B0 is GPU‑bound

62%

of B0’s 223 s is the predict stage (fp32 + 8‑pass TTA).

O6 is CPU‑bound

58%

once predict shrinks, preprocess dominates the 47 s.

Predict energy (mean)

14.7× less

B0 45.1 → O6 3.07 kJ/case; carbon 4.93 → 0.34 gCO₂. 4090, n=189.

Parallel producers measured

2.1×

3 workers preprocess at once: −30 % total energy, Dice‑identical. L40S, 10 cases.

What the ladder shows. The GPU‑side changes — TensorRT, fp16, less TTA — account for nearly all of the early savings, but only up to the point where the prediction bar drops below preprocess. After that the scan is CPU‑bound, and the remaining lever is to preprocess cases in parallel so the GPU stops waiting on data. Changes aimed at other stages — export‑side tricks, or torch.compile — make little difference here, because they optimise a step that is no longer the bottleneck.

What’s next — candidate config changes

The ladder ends at O6. These are the next rungs and what each one changes.

adopt
O6P — parallel workers. 2–3 CPU workers preprocess cases at once so the GPU never waits. Measured (L40S, 10 cases): 2.1× throughput, −30 % total energy, Dice‑identical. The one lever that beats O6’s CPU wall — ready to adopt on the 4090.
next
O7 — INT8 TensorRT. Quantize the engine fp16 → INT8 (post‑training, small calibration set, no retrain). Cheapest rung below O4 — if Dice holds. That’s the open question; calibration cases must be disjoint from eval.
tested
PyTorch forward flags. cudnn.benchmark + TF32. Free but marginal (~4 %). torch.compile was dropped — 2× slower on the sliding window.
research
Tiny‑net / combo cascade. A small net, or a tiny localizer that crops the kidney ROI before nnU‑Net. Attacks the CPU wall structurally; needs training — higher risk, higher reward.
research
GPU input‑resample. Move the preprocess resample onto the GPU. Shrinks the wall directly, but changes interpolation → breaks Dice‑comparability, so it’s its own study.

Mean per case over 189 held‑out cases (00400–00588), RTX 4090, NVML energy. Bars are stage split (preprocess / predict / export); totals are sum‑of‑stages. The O6 + parallel lever is measured (Modal L40S, 10 cases: −30 % total energy, 2.1× throughput, Dice‑identical); its dashed bar projects that effect onto this 4090 ladder. Preprocess is a single‑thread scipy resample.