Is there any example of Gradient Checkpointing to train large models with 3D dataset with limited GPU memory?

[faizan1234567]

Hi Author, thanks for the nice work and the help.

I'm working with a 3D medical imaging dataset, and when I increase the image resolution, I encounter memory issues that prevent the model from loading for training. This limitation hampers my ability to explore new ideas and more sophisticated architectural designs, especially given my compute and memory resource constraints.

I've learned about gradient checkpointing, which can help reduce GPU memory usage, and I see that it has already been implemented with the SwinUNETR model. Are there any easy-to-understand examples for non-sequential models? How can I implement gradient checkpointing in my custom model? Are there any guidelines or relevant documentation available in MONAI? I've searched for non-sequential model examples but haven't found anything intuitive online. Thank you in advance!