我正在尝试使用自定义数据集在 Jupyter 笔记本上训练 YOLOv5 模型。该数据集是一个口罩检测数据集,包含有/没有口罩的人的图像。我已经将注释转换为 YOLO 格式,并且我相信我已经编辑了所有必要的文件以反映类的数量(3:没有口罩、正确佩戴口罩、错误佩戴口罩)和培训/验证文件位置。
完成此操作后,我执行了以下命令:
!python train.py --img 256 --batch 8 --epochs 30 --data ./data/facemask.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --device 0
但我收到此错误:
Traceback (most recent call last):
File "./yolov5-master/train.py", line 404, in <module>
train(hyp)
File "./yolov5-master/train.py", line 79, in train
model = Model(opt.cfg).to(device)
File "/storage/facemask/yolov5-master/models/yolo.py", line 64, in __init__
m.stride = torch.tensor([64 / x.shape[-2] for x in self.forward(torch.zeros(1, ch, 64, 64))]) # forward
File "/storage/facemask/yolov5-master/models/yolo.py", line 91, in forward
return self.forward_once(x, profile) # single-scale inference, train
File "/storage/facemask/yolov5-master/models/yolo.py", line 108, in forward_once
x = m(x) # run
File "/opt/conda/envs/fastai/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/storage/facemask/yolov5-master/models/yolo.py", line 28, in forward
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
RuntimeError: shape '[1, 3, 8, 8, 8]' is invalid for input of size 8192
我发现 YOLOv3 的帖子提到应该更新 yolov3-spp.cfg 文件中的过滤器数量,但是我不相信 YOLOv5 有任何这样的文件。
有没有人有任何见解?
为了重现性,重新格式化的数据集和所有补充文件都可以在这里找到