search

WongKinYiu / yolor

implementation of paper - You Only Learn One Representation: Unified Network for Multiple Tasks (https://arxiv.org/abs/2105.04206)
GNU General Public License v3.0
1.99k stars 520 forks source link

Single class training #146

Open MrCri5py opened 2 years ago

MrCri5py commented 2 years ago

Hi, I am trying to train a model on a single class. However, when I run

"python train.py --batch-size 8 --img 640 640 --data data/custom.yaml --cfg cfg/custom_p6.cfg --weights yolor_p6.pt --device 0 --name custom_p6 --hyp hyp.scratch.1280.yaml --epochs 300 --single-cls --workers 1"

I get the following error message:

RuntimeError: shape '[3, 3, 6, 80, 80]' is invalid for input of size 4896000

I have edited the following lines in the cfg file: batch=8 subdivisions=6 width=640 height=640 classes=1

"every classes call was set to 1" What am I missing? Thanks!

raphychek commented 2 years ago

Hi! I had the same issue and followed this tutorial: https://github.com/ultralytics/yolov3/issues/102 Changing the filter values which were 255 to (5+n) * 3, so 18 in our case, fixed the issue for me.