使用YOLOv5模型训练： 显卡内存不够导致报错RuntimeError: CUDA out of memory

Logging results to runs\train\exp18
Starting training for 300 epochs...

     Epoch   gpu_mem       box       obj       cls    labels  img_size
  0%|          | 0/8 [00:03<?, ?it/s]
Traceback (most recent call last):
  File "E:\SelfLearning\Pytorch_Learning\yolov5\yolov5-6.0\train.py", line 620, in <module>
    main(opt)
  File "E:\SelfLearning\Pytorch_Learning\yolov5\yolov5-6.0\train.py", line 517, in main
    train(opt.hyp, opt, device, callbacks)
  File "E:\SelfLearning\Pytorch_Learning\yolov5\yolov5-6.0\train.py", line 315, in train
    pred = model(imgs)  # forward
  File "D:\DeepLearning\anzhuangCode\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "E:\SelfLearning\Pytorch_Learning\yolov5\yolov5-6.0\models\yolo.py", line 126, in forward
    return self._forward_once(x, profile, visualize)  # single-scale inference, train
  File "E:\SelfLearning\Pytorch_Learning\yolov5\yolov5-6.0\models\yolo.py", line 149, in _forward_once
    x = m(x)  # run
  File "D:\DeepLearning\anzhuangCode\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "E:\SelfLearning\Pytorch_Learning\yolov5\yolov5-6.0\models\common.py", line 45, in forward
    return self.act(self.bn(self.conv(x)))
  File "D:\DeepLearning\anzhuangCode\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "D:\DeepLearning\anzhuangCode\lib\site-packages\torch\nn\modules\activation.py", line 391, in forward
    return F.silu(input, inplace=self.inplace)
  File "D:\DeepLearning\anzhuangCode\lib\site-packages\torch\nn\functional.py", line 2048, in silu
    return torch._C._nn.silu(input)
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 2.00 GiB total capacity; 1.66 GiB already allocated; 0 bytes free; 1.72 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

进程已结束,退出代码1

遇到上述情况报错是由于模型batchsize设置的较大导致模型在训练时显卡内存不够。

遇到这种情况可以将batchsize减小。

在终端运行命令将batch修改合适的大小即可：
python train.py --img 640 --batch 4 --epochs 6