YOLOv5実行時のUnable to find a valid cuDNNエラーの原因と解決法

実行環境

OS: Ubuntu18.04
GPU: V100(16GB)

YOLOv5実行中に以下のエラーが発生

実行コード

python train.py --img 3200 --batch 16 --epochs 10 --data data.yaml --weights yolov5s.pt

エラー内容

  (...)
  File "C:\Program Files\Python38\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "C:\Program Files\Python38\lib\site-packages\torch\nn\modules\conv.py", line 423, in forward
    return self._conv_forward(input, self.weight)
  File "C:\Program Files\Python38\lib\site-packages\torch\nn\modules\conv.py", line 419, in _conv_forward
    return F.conv2d(input, weight, self.bias, self.stride,
RuntimeError: Unable to find a valid cuDNN algorithm to run convolution

原因と解決法

GPUのメモリ不足が原因のため、メモリ不足が起きない程度まで
batch sizeまたはimage sizeを修正する。