How to change cfg for custom dataset with different number of classes?

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 518 forks source link

How to change cfg for custom dataset with different number of classes? #34

Closed TheConstant3 closed 3 years ago

TheConstant3 commented 3 years ago

Darknet repo has guide "How to train (to detect your custom objects)", where described how change count filters and other parameters for custom dataset. If i change in yolor_p6.cfg number of classes to 2, i get error

  File "/workspace/yolor/models/models.py", line 400, in forward
    p = p.view(bs, self.na, self.no, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous()  # prediction
RuntimeError: shape '[4, 3, 7, 80, 80]' is invalid for input of size 6528000

WongKinYiu commented 3 years ago

https://github.com/WongKinYiu/yolor/issues/16

https://github.com/WongKinYiu/yolor/issues/25

TheConstant3 commented 3 years ago

Sorry, couldn't find it

Thank you so much!