open-mmlab / mmdetection

OpenMMLab Detection Toolbox and Benchmark
https://mmdetection.readthedocs.io
Apache License 2.0
29.67k stars 9.48k forks source link

Does Co-DETR supports automatic-mixed-precision training? #12000

Open taemmini opened 1 month ago

taemmini commented 1 month ago

When training using Co-DETR (using the co_dino_5scale_swin_l_16xb1_16e_o365tococo.py config) in an environment with 6 RTX4090 GPUs, MMCV==2.2.0, and MMdetection==3.3.0, the following error occurs:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.78 GiB. GPU 5 has a total capacity of 23.65 GiB of which 38.25 MiB is free. Process 4916 has 390.70 MiB memory in use. Including non-PyTorch memory, this process has 22.62 GiB memory in use. Of the allocated memory 16.82 GiB is allocated by PyTorch, and 5.21 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

While looking for solutions to this problem, I found the following section in tools/train.py: Copyparser.add_argument( '--amp', action='store_true', default=False, help='enable automatic-mixed-precision training')

I want to try training by activating this option. Is this possible with Co-DETR?