Repository navigation
YOLOv9 training runs out of GPU memory with a large max_labels #935
Description
I'm testing YOLO9 on the small object detection dataset "AI-TODv2", which can have as many as 2667 objects in a single 800x800 image. Training with max_labels: 3000 produces an OOM error:
OutOfMemoryError: CUDA out of memory. Tried to allocate 9.39 GiB. GPU 0 has a total capacity of 22.03 GiB of which 8.89 GiB is free. Including non-PyTorch memory, this process has 13.13 GiB memory in use. Of the allocated memory 12.86 GiB is allocated by PyTorch, and 46.79 MiB is reserved by PyTorch but unallocated.
I tested this on a L4 Colab GPU with LibreYOLO 1.5, Yolo9 tiny, an imgsz of 800, batch 16 and SGD.
Ultralytics yolo26n-p2 also has trouble, but it's capable of continuing:
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
1/100 21.7G 3.676 5.043 0.004033 889 800: 24% ━━╸───────── 175/701 1.4it/s 1:56<6:04WARNING ⚠️ CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU
1/100 20.3G 3.651 4.942 0.003974 420 800: 27% ━━━───────── 190/701 1.6it/s 3:27<5:16WARNING ⚠️ CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU
1/100 20.3G 3.61 4.77 0.003894 582 800: 31% ━━━╸──────── 219/701 1.9it/s 4:48<4:08WARNING ⚠️ CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU
1/100 20.3G 3.568 4.599 0.003802 396 800: 35% ━━━━──────── 251/701 1.7it/s 6:08<4:32WARNING ⚠️ CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU
1/100 20.3G 3.428 4.288 0.003597 522 800: 48% ━━━━━╸────── 340/701 1.8it/s 8:02<3:20WARNING ⚠️ CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU
1/100 21.6G 3.398 4.23 0.003563 206 800: 51% ━━━━━━────── 360/701 2.1it/s 9:35<2:41WARNING ⚠️ CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU
1/100 20.4G 3.296 3.99 0.003391 494 800: 65% ━━━━━━━╸──── 456/701 1.8it/s 11:38<2:18WARNING ⚠️ CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU
Activity
Merge pull request #939 from LibreYOLO/935-yolo9-per-image-assignment
EHxuban11 commentedon Oct 17, 2026
Hello @ManuelZ, thank you for reporting this!
This was a bug in the library: with
max_labels: 3000, for example, every image in the batch reserved memory for 3000 objects, even if it had far fewer. This has been fixed on dev in #939.If you still hit an out-of-memory error after this, it means the GPU simply doesn't have enough memory, and there isn't much we can do about that. You could either train directly on the CPU (if you have more RAM than your GPU has VRAM), although it will be very slow, or decrease the batch size. If you want a smaller batch size without the spikiness in training that small batches cause, you can use
nbs(the nominal batch size) to accumulate gradients: for example,batch=4, nbs=16trains with an effective batch size of 16.Let me know how it goes with AI-TODv2! By the way, what is this dataset? Is it famous? Could you tell me more about it? Thank you!
ManuelZ commentedon Oct 19, 2026
Hi @EHxuban11, I'm testing it out and it's working great! I'm not getting OOM errors anymore with the same setup, thank you!
Regarding AI-TODv2, I don't know if it's famous, but I chose it because it's a big dataset, composed of various smaller ones.
The following is taken from the AI-TOD paper. As far as I know, AI-TODv2 is just AI-TOD with better annotations.
We build the AI-TOD based on the publicly available large-scale aerial image datasets:
The details of these datasets are as follows:
DOTA-v1.5 trainval set: It is an upgraded version of original DOTA-v1.0 dataset, and has been employed for performance evaluation in Detecting Objects in Aerial Images (DOAI2019)2. DOTA-v1.5 trainval set contains 1,869 images ranging in size from 800 × 800 to 4000 × 4000 pixels and 280,196 object instances which are labeled into 16 classes (e.g., ship, small-vehicle, storage-tank).
xView training set: It is a large-scale object detection dataset which consists of 1,415km2 of WorldView-3 at 30cm resolution. The annotated dataset for object detection contains over 1 million object instances across 60 classes, including various types of vehicles, planes, and ships [19].
VisDrone2018-Det trainval set: It consists of 7,019 images captured by drone platforms in different places at different heights [20]. Images are manually annotated with bounding boxes and 10 predefined classes (e.g., pedestrian, person, car).
Airbus-Ship trainval set: It is a dataset for ship detection in Kaggle challenge. Airbus-ship trainval set consists of 42,559 images and 81,724 ships, all objects are annotated by polygon.
DIOR trainval+test set: It includes 23,463 images and 192,472 object instances covered by 20 classes (e.g., airplane, ship, windmill).
Match YOLO9 labels one image at a time
EHxuban11 commentedon Oct 21, 2026
Hello @ManuelZ , thanks for checking, happy that it works now.
AI TOD v2 could be used to fine tune object detectors specialized on small objects in the future in LibreYOLO.. interesting.