ControlNet training out of memory

[MaybeRichard]

I'm following the instruction in the README.md and try to train ControlNet with KiTS dataset. I've met the following problem and I can't find solution from the instruction1.

1. running on RTX 4090 with command python -m scripts.train_controlnet -c ./configs/config_maisi.json -t ./configs/config_maisi_controlnet_train.json -e ./configs/environment_maisi_controlnet_train.json -g 1, and the problem is " CUDA out of memory." I check the 'config_maisi_controlnet_train.json' file and the batch_size is 1. I also check the C4KC-KiTS_subset.json and I found the dimention is 512512512. I think this is because the input size is too large but the 'C4KC-KiTS_subset.json' and the correspoding dataset is downloaded from CLI, I'm nor sure could I change the dimention size or do I need to regenerate the json file. I also wanna know how to generate my own dataset that met the format requirement like 'C4KC-KiTS_subset.json'.

[KumoLiu]

Hi @MaybeRichard, you can try to modify the num_splits parameter in AutoencoderKlMaisi. This adjustment reduces the spatial splitting of feature maps in convolutions, which is particularly beneficial given the smaller input size. (This change helps convert convolutions to a for-loop based approach, thereby conserving GPU memory resources.)

For more details, please refer to:
https://github.com/Project-MONAI/tutorials/tree/main/generation/maisi#time-cost-and-gpu-memory-usage

[MaybeRichard]

Thanks, I'ii try it.