How much GPU memory does the training need？ about yolos HOT 2 CLOSED

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:
  

  YOLOS/models/backbone.py

  Lines 325 to 326 in 2e10dc4

  	if self.use_checkpoint:
  	x = checkpoint.checkpoint(self.blocks[i], x) # saves mem, takes time

  
  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.

from yolos.

Thanks for your reply. I will try gradient checkpointing~

from yolos.