Почему я не могу обучить YOLOv5 на моем наборе данных?

Question

Я пытаюсь обучить модель YOLOv5 на записной книжке Jupyter, используя настраиваемый набор данных. Набор данных представляет собой набор данных обнаружения маски лица, содержащий изображения людей с масками или без них. Я преобразовал аннотации в формат 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 имеет такой файл.

У кого-нибудь есть идеи?

Для воспроизводимости здесь доступны переформатированный набор данных и все дополнительные файлы.

Answer 1

попробуй добавить x[i]=self.m[i](x[i])строка 28 в yolo.py. Вот так

      def forward(self, x):
   # x = x.copy()  # for profiling
    z = []  # inference output
    self.training |= self.export
    for i in range(self.nl):
        #[1, 128, 80, 80]
        bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
        x[i] =self.m[i](x[i]) # **add this code here**
        x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

если это не работает, возможно, что-то не так с вашими весами, попробуйте использовать файлы *.pt, которые соответствуют версии yolov5.