4GB GPU运行Detectron2时CUDA内存不足问题求助

错误信息 #

RuntimeError: CUDA out of memory. Tried to allocate 54.00 MiB (GPU 0;
4.00 GiB total capacity; 624.92 MiB already allocated; 2.02 GiB free; 720.00 MiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.

你的配置代码 #

cfg = get_cfg()
cfg.MODEL.DEVICE = "cuda"

cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("pan_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 300    # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.SOLVER.STEPS = []         # do not decay learning rate
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 11  # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
# NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

具体解决方案 #

针对4GB显存的GPU，按优先级调整以下配置：

1. 降低批量大小
   将cfg.SOLVER.IMS_PER_BATCH从2改为1，这是最直接减少显存占用的方式：

   cfg.SOLVER.IMS_PER_BATCH = 1
   

2. 开启混合精度训练
   启用自动混合精度（AMP），可大幅降低显存占用（约50%），同时不影响训练精度：

   cfg.SOLVER.AMP.ENABLED = True
   

3. 减小ROI批量大小
   进一步降低每个图像的ROI样本数，将cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE从128调整为64或32：

   cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
   

4. 降低输入图像分辨率
   缩小训练时的图像尺寸，减少张量占用的显存：

   cfg.INPUT.MIN_SIZE_TRAIN = (400,)
   cfg.INPUT.MAX_SIZE_TRAIN = 600
   

5. 优化显存碎片问题
   按照错误提示设置显存分裂阈值，在初始化cfg后添加：

   import torch
   torch.cuda.max_split_size_mb = 128
   

6. 关闭CUDA基准测试
   如果显存碎片问题仍存在，关闭cudnn基准测试，避免预分配过多显存：

   torch.backends.cudnn.benchmark = False
   

建议先尝试前3种方法，通常就能解决4GB显存的不足问题。如果仍有问题，再逐步添加后面的优化项。

内容的提问来源于stack exchange，提问作者Kiran Roy