2

I am new to deep learning. I am doing a school project where I am trying to train a YOLOv5 model on VisDrone dataset. My training set has 4911 images and validation set has more than 3000 images but less than 4000. I am using google Colab pro. As far as I know it has 32gb ram capacity and GPU VRAM=15-16 GB. If I let the model load data automatically it's showing that "cuda out of memory". What strategy can I take to slove this problem? Should I customize the dataloader which is dataloaders.py file. How do I do that?

ShehjadK
  • 31
  • 7

1 Answers1

1

Usually cuda out of memory occurs due to batch size, it is much better if you let the trainer decide the batchsize it self. To do this replace line number 442 in train.py from:

 parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')

to this:

 parser.add_argument('--batch-size', type=int, default=-1, help='total batch size for all GPUs, -1 for autobatch')

This should solve the out of memory issue.

Azhan Mohammed
  • 340
  • 2
  • 8
  • Now I am getting this error. – ShehjadK Aug 14 '22 at 04:00
  • AutoBatch: Computing optimal batch size for --imgsz 4928 AutoBatch: CUDA:0 (Tesla P100-PCIE-16GB) 15.90G total, 0.36G reserved, 0.35G allocated, 15.19G free tcmalloc: large alloc 2331377664 bytes == 0x10f27e000 @ 0x7f606ed99b6b – ShehjadK Aug 14 '22 at 04:02
  • Why are you using such a large image size? Try taking small crops of the image rather than using the complete image at once. While working with drone/satellite images, I usually crop out the image into patches of 600x600 pixel. – Azhan Mohammed Aug 14 '22 at 17:41
  • I understood the problem. Thanks – ShehjadK Aug 15 '22 at 10:47