Debug performance

[constantinpape]

Trying to reproduce / debug the performance issues reported by @SarahMuth and @lufre1 .

[constantinpape]

Vesicle Segmentation

Trying to debug the problems with vesicle segmentation performance reported by @SarahMuth .
I can't really reproduce this, the segmentation times look normal for me, see below.

To reproduce:

1. Get interactive session (without any qualifiers):

srun -p grete:shared -G A100:1 --pty bash

2. Run segmentation via CLI:

synapse_net.run_segmentation -i /mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/original_imod_data/20240909_cp_datatransfer/10_tem_single_release/TEM_250nm_Culture5_SingleRelease/ -o out -m vesicles_3d

The segmentation runs in ca. 36 minutes (= ca. 2:15 min / tomogram):

16/16 [36:23<00:00, 136.47s/it]

For a single tomogram, the individual steps take this long:

• Rescaling: 9.752075910568237 s
• Prediction: 25.630613803863525 s
• Distance Transform: 44.692020654678345 s
• Connected components: 2.0873935222625732 s
• Distance Transform: 39.08389163017273 s
• Watershed: 14.07514500617981 s
• Size filter: 2.796635866165161 s
• Resize: 3.0298449993133545 s

[constantinpape]

Mito Segmentation

Trying to debug the problems / out of memory error in mito segmentation reported by @lufre1 .
I can't reproduce this either, everything runs as expected:

To reproduce:

1. Get interactive session (without any qualifiers):

srun -p grete:shared -G A100:1 --pty bash

2. Run segmentation via CLI

synapse_net.run_segmentation -i /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250212_test_I_h5_s2 -o out -m mitochondria --tile_shape 48 512 512 --halo 4 128 128 --verbose -s raw --data_ext .h5

The segmentation runs in ca. 6:30 minutes (= ca 1:10 min / tomogram) without further issues:

5/5 [06:22<00:00, 76.40s/it]

For an individual tomogram, the individual steps take this long:

• Predictio: 46.622530698776245 s
• Distance transform: 7.458238124847412 s
• Watershed: 8.343072175979614 s
• Size filter in 1.1268229484558105 s

(Note that the h5s don't contain voxel sizes, so the resizing step is skipped)

[lufre1]

Mito Segmentation

@constantinpape The problem occurs only with .mrc files. The .h5 files are downscaled by a factor of two and i created those to cirumvent the problem.
Prior to this test I recreated my mamba env.

what i used:

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 --mem 32G /bin/bash

and then (pulled some mrc files from fidi we can use)

(synapse) [u12103@ggpu135 synapse-net]$ synapse_net.run_segmentation -i /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250228_for_Mito_Seg/ -o out -m mitochondria --tile_shape 48 512 512 --halo 4 128 128 --verbose

and that is the output:

Determined tile size for CUDA: {'tile': {'x': 512, 'y': 512, 'z': 64}, 'halo': {'x': 64, 'y': 64, 'z': 16}}
Processing files:   0%|                                                                                                                                                                                                                               | 0/13 [00:00<?, ?it/s)
Segmenting mitochondria in volume of shape (440, 1981, 1854)
Rescaled volume from (440, 1981, 1854) to (185, 832, 779) in 24.4645357131958 s
predict with halo: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 80/80 [00:33<00:00,  2.38it/s]
Prediction time in 33.67737650871277 s███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 80/80 [00:33<00:00,  2.47it/s]
Compute distance transform: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:07<00:00,  4.31it/s]
Compute distance transform in 8.867842197418213 s██████████████████████▍                                                                                                                                                                      | 6/32 [00:07<00:21,  1.19it/s]
Label all sub-blocks: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 1060.03it/s]
Merge labels across block faces: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 1469.91it/s]
Write blocks: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 808.22it/s]
Compute connected components in 0.6240131855010986 s                                                                                                                                                                                                  | 0/32 [00:00<?, ?it/s]
Seeded watershed: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:02<00:00, 15.60it/s]
Compute watershed in 2.917395830154419 s███████████████████████▎                                                                                                                                                                              | 6/32 [00:01<00:06,  3.76it/s]
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 205.74it/s]
Size filter in 0.6509954929351807 s                                                                                                                                                                                                           | 2/32 [00:00<00:01, 18.10it/s]
Resized prediction back to original shape (440, 1981, 1854) in 16.83793020248413 s
Saved segmentation to out/20250228_for_Mito_Seg/36194_B4_66K_TS_R01A_SC_01_rec_crop2_prediction.tif.
Processing files:   8%|████████████████▍                                                                                                                                                                                                     | 1/13 [01:58<23:36, 118.05s/it)
Segmenting mitochondria in volume of shape (471, 1779, 1779)
Killed

EDIT

with no mem specified:

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 /bin/bash

the output:

(synapse) [u12103@ggpu131 synapse]$ synapse_net.run_segmentation -i /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250228_for_Mito_Seg/ -o out -m mitochondria --tile_shape 48 512 512 --halo 4 128 128 --verbose
Determined tile size for CUDA: {'tile': {'x': 512, 'y': 512, 'z': 64}, 'halo': {'x': 64, 'y': 64, 'z': 16}}
Processing files:   0%|                                                                                                                                                                                                                               | 0/13 [00:00<?, ?it/s)
Segmenting mitochondria in volume of shape (504, 1831, 1868)
Rescaled volume from (504, 1831, 1868) to (212, 769, 785) in 26.210493564605713 s
predict with halo: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 96/96 [00:40<00:00,  2.39it/s]
Prediction time in 40.23682236671448 s███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 96/96 [00:40<00:00,  2.47it/s]
Compute distance transform: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:06<00:00,  4.60it/s]
Compute distance transform in 8.10959506034851 s███████████████████████████████████████████████████████▏                                                                                                                                     | 11/32 [00:06<00:07,  2.80it/s]
Label all sub-blocks: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 766.17it/s]
Merge labels across block faces: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 4355.74it/s]
Write blocks: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 344.69it/s]
Compute connected components in 0.6494076251983643 s                                                                                                                                                                                                  | 0/32 [00:00<?, ?it/s]
Seeded watershed: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:01<00:00, 27.01it/s]
Compute watershed in 2.003214120864868 s                                                                                                                                                                                                      | 2/32 [00:01<00:15,  1.95it/s]
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 198.47it/s]
Size filter in 0.6730284690856934 s██████▍                                                                                                                                                                                                    | 5/32 [00:00<00:00, 31.07it/s]
Killed

and with 128GB mem specified:

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 --mem 128 /bin/bash

output:

(synapse) [u12103@ggpu131 synapse]$ synapse_net.run_segmentation -i /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250228_for_Mito_Seg/ -o out -m mitochondria --tile_shape 48 512 512 --halo 4 128 128 --verbose
Killed

and with 256 mem specified:

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 --mem 256 /bin/bash

the ouput:

(synapse) [u12103@ggpu131 synapse]$ synapse_net.run_segmentation -i /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250228_for_Mito_Seg/ -o out -m mitochondria --tile_shape 48 512 512 --halo 4 128 128 --verbose
Killed

there might still be issues with the environemnt - i will have a look into that...

[constantinpape]

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 --mem 32G /bin/bash

Why are you setting it to 32G ? That's tiny. Either set to at least 128G (or better 256G) or don't constraint it (not sure what the default is, but for me everything worked with it.)

[lufre1]

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 --mem 32G /bin/bash

Why are you setting it to 32G ? That's tiny. Either set to at least 128G (or better 256G) or don't constraint it (not sure what the default is, but for me everything worked with it.)

i tested with no mem specified:

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 /bin/bash

then i found 30GB of RAM:

(synapse-net) [u12103@ggpu131 synapse-net]$ scontrol show job $SLURM_JOB_ID | grep TRES
   ReqTRES=cpu=8,mem=30560M,node=1,billing=300,gres/gpu=1,gres/gpu:a100=1
   AllocTRES=cpu=8,mem=30560M,node=1,billing=300,gres/gpu=1,gres/gpu:a100=1

and got:

(synapse-net) [u12103@ggpu131 synapse-net]$ synapse_net.run_segmentation -i /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250228_for_Mito_Seg/ -o out -m mitochondria --tile_shape 48 512 512 --halo 4 128 128 --verbose
Determined tile size for CUDA: {'tile': {'x': 512, 'y': 512, 'z': 64}, 'halo': {'x': 64, 'y': 64, 'z': 16}}
Processing files:   0%|                                                                                                                                                                                                                               | 0/13 [00:00<?, ?it/s]Rescaling the data at /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250228_for_Mito_Seg/20250228_for_Mito_Seg/36194_B4_66K_TS_SC_08_rec_crop.mrc by (0.42000001179423313, 0.42000001179423313, 0.42000001179423313) to match the training voxel size (2.07, 2.07, 2.07)
Segmenting mitochondria in volume of shape (471, 1779, 1779)
Rescaled volume from (471, 1779, 1779) to (198, 747, 747) in 23.79210615158081 s
predict with halo: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 45/45 [00:19<00:00,  2.28it/s]
Prediction time in 19.77400779724121 s███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 45/45 [00:19<00:00,  2.45it/s]
Compute distance transform: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18/18 [00:07<00:00,  2.31it/s]
Compute distance transform in 8.241600751876831 s████████████████████████████████████████▉                                                                                                                                                    | 5/18 [00:07<00:13,  1.01s/it]
Label all sub-blocks: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 479.02it/s]
Merge labels across block faces: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 1620.43it/s]
Write blocks: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 167.58it/s]
Compute connected components in 0.6296136379241943 s███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉                        | 16/18 [00:00<00:00, 149.36it/s]
Seeded watershed: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18/18 [00:01<00:00, 15.47it/s]
Compute watershed in 1.7455918788909912 s█████████████████▊                                                                                                                                                                                   | 3/18 [00:01<00:04,  3.06it/s]
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 105.12it/s]
Size filter in 0.6214752197265625 s                                                                                                                                                                                                           | 1/18 [00:00<00:02,  7.17it/s]
Resized prediction back to original shape (471, 1779, 1779) in 15.575767755508423 s
Saved segmentation to out/20250228_for_Mito_Seg/36194_B4_66K_TS_SC_08_rec_crop_prediction.tif.
Processing files:  15%|█████████████████████████████████                                                                                                                                                                                      | 2/13 [01:31<08:23, 45.75s/it]Rescaling the data at /scratch-grete/projects/nim00007/data/mitochondria/cooper/20250228_for_Mito_Seg/20250228_for_Mito_Seg/36194_B4_66K_TS_SC_19_rec_crop2.mrc by (0.42000005786545613, 0.42000001179423313, 0.42000001179423313) to match the training voxel size (2.07, 2.07, 2.07)
Segmenting mitochondria in volume of shape (504, 1831, 1868)
Killed

[constantinpape]

i tested with no mem specified:

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 /bin/bash

then i found 30GB of RAM:

This is because you specify -c 8. It apparently limits the memory then compared to not asking for any constraints. You can see my exact commands here: #108 (comment)

But you can also just ask for more memory (128G or 256G) and the problem should be solved.

[lufre1]

i tested with no mem specified:

srun --pty -p grete:shared -G A100:1 -c 8 -t 2-00:00:00 /bin/bash

then i found 30GB of RAM:

This is because you specify -c 8. It apparently limits the memory then compared to not asking for any constraints. You can see my exact commands here: #108 (comment)

But you can also just ask for more memory (128G or 256G) and the problem should be solved.

with this srun command i get:

[u12103@glogin9 synapse]$ srun --pty -p grete:shared -G A100:1 /bin/bash 

[u12103@ggpu104 synapse]$ micromamba activate /scratch-grete/usr/nimlufre/envs/synapse-net
(synapse-net) [u12103@ggpu104 synapse]$ scontrol show job $SLURM_JOB_ID | grep TRES
   ReqTRES=cpu=1,mem=3820M,node=1,billing=300,gres/gpu=1,gres/gpu:a100=1
   AllocTRES=cpu=64,mem=244480M,node=1,billing=300,gres/gpu=1,gres/gpu:a100=1

ok, with this configuration I do not run into any issues

[constantinpape]

Thanks for reporting @lufre1 . So it's consistent that this issue is due to memory. I will go ahead and merge this issue, as we don't need to address anything else. If you find a volume where we run into issues even with 256GB let me know, and then we can check how to fix memory for it.