How to train model using mixed precision fp16?

[Bhavay-2001]

Hello author, I tried using the original DiT model for training but facing out of memory issue. I saw your repository which implements DiT using memory constraints. In the README file, I saw you used a mixed_precision argument but I couldn't find it anywhere in the code. I just want to copy the model architecture file and adjust it according to my implementation of the work. Can you please tell which model arch uses less memory constraints as it is a bit confusing to me to understand so just clarifying.

[chuanyangjin]

Hello @Bhavay192,

The "--mixed_precision" argument is automatically taken care of by the accelerate library. Therefore, there's no direct reference to it in the code.

The gradient checkpointing significantly reduces memory usage. You can find the relevant implementation in the models.py, lines 233-237 and line 251.

[Bhavay-2001]

Heyy, thanks for the reply.

I'm trying your model implementation on medical image anomaly detection on Brats21 dataset. Currently the input image size is 96, and patch_size(another param) is 48. If I run with these configs, I'm facing out of memory issue.

However, If I reduce the input image and patch size, the code runs well for few epochs until it crashes to throw an error produced by the changed patch_size. I believe the code will run perfect with default settings(image_size - 96, patch_size-46). How can I convert the model to fp16, any ideas?

Thanks