How much GPU memory does the training need？

[LuckyDC]

How much GPU memory does the training need？ It always raise out-of-memory error for me.

[Yuxin-CV]

Hi @LuckyDC, thanks for your interest in YOLOS.

We train all YOLOS models on 1 node with 8 x 11G 2080Ti GPUs.

• For the YOLOS-Ti, the batch_size_per_GPU can set to 2 or 3, we use 2. 150 epochs training takes ~4 days.
• For YOLOS-Small and YOLOS-Base, we use a hack called gradient checkpointing: https://github.com/hustvl/YOLOS/blob/2e10dc45b5effc015bd64d597b546caf47fb3c0e/models/backbone.py#L325-L326 A tutorial can be found here. This flag can save a lot of GPU memory, at a cost of more training time: 150 epochs training takes ~10 days for YOLOS-Small, and takes ~20 days for YOLOS-Base.

Recently, we have access to an 8 x 24G 3090 GPUs node, we find it takes ~6 days to train 150 epochs for YOLOS-Small.

[LuckyDC]

Thanks for your reply. I will try gradient checkpointing~