[constantinpape]

No description provided.

[constantinpape]

@lufre1 @stmartineau99 : this uses a new implementation of the prediction function from torch_em that supports batched prediction and that also uses sub-processes with a queue to fetch and preprocess input blocks and write output blocks. This avoids that the GPU idles while waiting for I/O and preprocessing, so should speed up the inference.

Please check out once that this works for your current applications, then we can merge it.

[lufre1]

I checked #172 (51f7310) on our mitochondria / cristae pipeline. The setup: mitochondria2 and cristae5 (cristae with the mito mask) on 4 tomograms from Yerin's data; 3 tilings (32×256×256, the new default from #188, and 64×512×512); an A100 MIG 1g.10gb slice, an A100-40GB and an A100-80GB. The old predict_with_halo and the PR ran in the same process on the same input, with 2 repeats after a warm-up. The timed part is get_prediction_torch_em.

A bug: predict_with_halo_pipelined deadlocks instead of raising

When the forward pass raises (in our case CUDA OOM), the prediction can hang forever instead of raising the error. In one run, cristae with 64×512×512 tiles at batch_size=4 on the 40 GB A100 sat for 7 h without writing a block, until it hit the job's time limit.

Cause (torch_em 0.10.8, util/prediction.py): in the producers' finally, the last producer signals the consumers with a blocking input_queue.put(_STOP). That put has no timeout and doesn't check stop_event.

• If all blocks are already queued and the queue is full, that producer blocks until the consumer takes an item.
• If the consumer then fails, it exits, nothing drains the queue, and the main thread waits in join() forever.

This happens exactly when the failing batch is the one taken with batch_size + queue_size (= 3 × batch_size on one device) blocks left. With a mask, only the non-skipped blocks count; our masked cristae volume had 12 blocks to predict at bs4. It also hits the default batch_size=1 with 3 blocks. The consumer's output_queue.put(_STOP) has the same pattern.

A CPU repro, with a model that raises in forward:

11 blocks, batch size 4 (queue 8): raised RuntimeError
12 blocks, batch size 4 (queue 8): HANGS
13 blocks, batch size 4 (queue 8): raised RuntimeError
 2 blocks, batch size 1 (queue 2): raised RuntimeError
 3 blocks, batch size 1 (queue 2): HANGS
 4 blocks, batch size 1 (queue 2): raised RuntimeError

Using _safe_put for both STOP signals fixes it: all cases raise. With the fix, outputs are unchanged and still identical to predict_with_halo on the success path.

            if remaining_producers.decrement() == 0 and not stop_event.is_set():
                try:
                    for _ in range(n_consumers):
                        _safe_put(input_queue, _STOP, stop_event)
                except _Aborted:
                    pass

(and the same for output_queue / num_write_workers in the consumer.) The repro script is repro_pipelined_deadlock.py, which I can share.

Correctness

• batch_size=1 gives bit-identical results to predict_with_halo in all 38 settings: 3 GPUs, both models, 3 tilings, with and without mask, plus a lazy ResizedVolume input with an n5 output (the scalable_segmentation path) and dtype=torch.bfloat16. It is also deterministic. On the 1g.10gb slice it also reproduces our cached production predictions (made with the old code) bit for bit.
• batch_size>1 is never bit-identical: the max abs difference is up to 5e-3 in fp32 and 8e-3 in bf16. In 11 of 79 settings it is also not reproducible run to run, probably because which blocks end up in the same batch depends on thread timing.
  • The effect on our segmentations is small: across 80 comparisons the number of objects never changes, and all objects match with IoU ≥ 0.988. The lowest is one mito object, which changes by 7k voxels.
• Unit tests: same result on the PR and on main (ffad121), 168 passed and 9 skipped. The 2 errors are import errors in test/training on both: get_unetr_model is missing in torch_em 0.10.8.

Speed (sum over the tomograms, speed-up vs. old)

GPU	case	bs1	bs2	bs4	bs8
1g.10gb	mito 32×256×256 (4 tomograms)	1.00	1.01	0.66	–
1g.10gb	cristae 32×256×256 (4)	1.01	0.99	0.46	–
A100-40GB	mito 32×256×256 (4)	1.00	1.15	1.18	1.18
A100-40GB	mito default from #188 (4)	0.98	1.10	1.13	1.11
A100-40GB	mito 64×512×512 (4)	0.91	0.16	OOM	–
A100-40GB	cristae 64×512×512 (1)	0.99	0.08	hangs (above)	–
A100-80GB	cristae 32×256×256 (3)	1.02	0.99	0.98	0.98
A100-80GB	cristae 64×512×512 (3)	0.93	0.93	–	–
1g.10gb	lazy ResizedVolume → n5 (scalable path), mito	1.6	1.5	1.0	–

• On the small slice the GPU is already 97–99% busy with the old code, so pipelining gains nothing.
• The clear win is a lazy input: the loader threads parallelize the per-block resize.
• Batching only helps mito with small tiles on a large GPU (up to 1.18×), and it never helps the masked cristae prediction.
• With 64×512×512 tiles, bs1 is 7–9% slower than the old code (mito on 40 GB, cristae on 80 GB).
• Close to the memory limit, batching gets much slower: bs2 with 64×512×512 tiles on 40 GB is 6× (mito, 36 GB peak) to 13× (cristae) slower. The log fills with CUDACachingAllocator OOM warnings, so cuDNN apparently falls back to algorithms without workspace. On 80 GB the same bs2 runs at bs1 speed.

From our side

• The deadlock should be fixed first (in torch_em). It can also hit the default batch_size=1, and every synapse-net prediction now goes through this function.
• After that, fine to merge with batch_size=1 as the default: bit-identical to predict_with_halo.
• I'd suggest the docstring / --batch_size help says that values > 1 change the results slightly and are not reproducible run to run. It should also say that batches close to the GPU memory limit can be much slower than bs1.

[constantinpape]

I have fixed the error regarding the non-terminating prediction in constantinpape/torch-em#901 and also made a release of torch_em that includes the fix.

[lufre1]

Thanks! I re-checked #172 (51f7310) with torch_em 0.10.10 on the A100 MIG 1g.10gb slice, with the same scripts as before:

• The deadlock is fixed.
  • The CPU repro now raises in all cases, including the two that hung (12 blocks at bs4, 3 at bs1). A control with 0.10.8 still hangs there.
  • The real case also raises now: cristae, 64×512×512 tiles, batch_size=4, 12 non-masked blocks, OOM on the first batch. 0.10.10 raises the CUDA OOM after ~6 s; with 0.10.8 the same run hangs until killed.
• No regression.
  • batch_size=1 is still bit-identical to predict_with_halo, for mito and cristae on 2 tomograms with 2 repeats. The predictions also have the same hashes as with 0.10.8.
  • Speed bs1/old is 1.00–1.03.
  • The lazy ResizedVolume → n5 path is still bit-identical at bs1 and 1.6× faster.
• Unit tests: all pass (187 run, 10 skipped). A correction to my earlier comment: the 2 test/training import errors came from my env having torch_em 0.10.8, below the >=0.10.9 this repo requires. They are gone with ≥0.10.9.

One suggestion: bump the requirement to torch_em>=0.10.10 in setup.py and environment.yaml in this PR. With 0.10.8 or 0.10.9 any error in the forward pass can hang the prediction, also at the default batch_size=1. 0.10.10 is already on conda-forge, so CI can pick it up right away.

From our side this is good to merge.