How to customize data loader of yolov5 to train on VisDrone dataset?

Question

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?

Answer 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.