训练yolov5导致CUDNN_STATUS_NOT_INITIALIZED错误

Training yolov5 causes a CUDNN_STATUS_NOT_INITIALIZED error

我正在关注 this guide,没有做任何更改。我正在使用带有深度学习 ami 的 aws 服务器:Deep Learning AMI (Ubuntu 18.04) 版本 40.0

我尝试将我的自定义数据集更改为 coco 数据集和自定义数据集的一小部分。 批量大小似乎无关紧要,CUDA 和其他驱动程序似乎可以工作。

批处理开始训练过程时抛出异常。这是完整的堆栈跟踪:

Logging results to runs/train/exp66
Starting training for 5 epochs...

     Epoch   gpu_mem       box       obj       cls     total   targets  img_size
  0%|                                                                                                                                                                                                                 | 0/22 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "train.py", line 533, in <module>
    train(hyp, opt, device, tb_writer, wandb)
  File "train.py", line 298, in train
    pred = model(imgs)  # forward
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/yolov5/models/yolo.py", line 121, in forward
    return self.forward_once(x, profile)  # single-scale inference, train
  File "/home/ubuntu/yolov5/models/yolo.py", line 137, in forward_once
    x = m(x)  # run
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/yolov5/models/common.py", line 113, in forward
    return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/yolov5/models/common.py", line 38, in forward
    return self.act(self.bn(self.conv(x)))
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 399, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 395, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED

我使用 conda 修复了它,我克隆了图像附带的 pytorch 环境,它运行良好。虽然我仍然不知道原因。

我不知道为什么,但似乎 torch 1.8 是建立在旧版本的 cuda 上的。 此外,由于 pytorch 有自己的 cuda,它似乎并不关心你的机器上有什么版本。 更改 torch 版本(并匹配兼容的 tochvision)解决了我的问题。

在我的例子中,我做了如下操作:

  1. 更改了“requirements.txt”中的两行:

手电筒==1.7.1

torchvision==0.8.2

  1. 使用 python=3.8
  2. 创建了全新的 conda 环境
  3. 激活环境
  4. 已安装来自修改文件的要求:

$ pip install -r requirements.txt

希望对大家有所帮助:)

我 运行 在尝试在脚本中训练 yolov5 时遇到了类似的问题。我发现升级到 torch==1.9.0 和 torchvision==0.10.0 也可以(如果你不想如上所述降级)