Providing anatomical labels as input

[BennoGesierich]

Is it possible to provide anatomical labels as input modality?

I tried to do this by adding a map with the anatomical labels as input modality, specifying 'noNorm' for it in the dataset.json, to avoid normalization. In addition, I modified my custom trainer class and the get_moreDA_augmentation function, to allow a similar treatment of this label map during loading and data augmentation as that applied to the images from previous training stages by the cascade trainer. In particularly, I tried to move the label map in a customized version of generate_train_batch to seg in the data dictionary and to move it back to data at the end of the data augmentation using the MoveSegAsOneHotToData class. It looks like it is working in principle, but the training gets always stuck in an early phase, but even not always in the same phase. It can be even different from fold to fold. Sometimes it gets stuck during _ = self.tr_gen.next() or self.val_gen.next(), and sometimes only afterwards, or even after the first few iterations.
Therefore, I suspect that the problem is somehow related to a performance issue, even so I am working with a GeForce RTX 3090 (VRAM 24GB) and also RAM should be sufficient.

Hopefully, not to complex for an answer. In any case, I would be very grateful for any hint pointing me in the right direction.

[FabianIsensee]

Hey so what concretely do you mean by 'getting stuck'? Does it never progress? Is nnU-Net still using the CPU/GPU when it is 'stuck'?
The cascade augmentations can be really CPU intensive, so it could be that you are just stuck in a severe CPU bottleneck. The larger you define the kernel sizes (the amount of erosion, dilation etc) the longer the computation will take. So take a close, careful look at your parametrization as well as your actual resource usage (CPU, GPU and RAM) during training

[BennoGesierich]

Hi Fabian, thank you for your reply!
It might indeed be the case that it gets stuck during data augmentation or batch generation. However, it seems like it has still plenty of resources left. Also, the cascade augmentation usually runs on our hardware.

I ran again two tests, monitoring the resources. In the fist attempt with fold=0 it got stuck already when executing the command _ = self.tr_gen.next() in the network_trainer.py run_trainer function. Is this command already preparing the fist batch for training or is this still part of setting the batch generator up? In the second attempt with fold=1 it got stuck only later after the first iteration. But, when repeating the test with the fold=1, it also got stuck already during _ = self.tr_gen.next() in the run_trainer function. I inserted some print_to_log_file commands in the code to verify were it got stuck.

However, the resources are far from exhausted.
In the first test it used after being blocked:
11.5 GiB RAM (2.3% of total RAM)
4.5% CPU usage
1001 MiB VRAM (of total 48 GiB)
0% GPU usage
In the second test it used after being blocked:
22.5 GiB RAM (4.5% of total RAM)
4.4% CPU usage
4.38 GiB (of total 48 GiB)
0% GPU usage

I know, my question is not easy to answer, given that I modified the code. But if you have any idea, which parameters I have to check, because they could cause the data loader and batch generator to be blocked this way, I would be very grateful.