gau-nernst/vision-toolbox

csd = ckpt['model'].float().state_dict()

Closed this issue · 2 comments

hey,dude.I find you are the only one to train imagenet in yolov5.
I wanna to use your pretrained model in yolov5 for object detecting.
But it cann't take it in raw yolov5?

python train.py --data tinyObject.yaml --cfg ./models/yolov5m.yaml  --weights darknet_yolov5m-a1eb31bd.pt

error:  File "train.py", line 123, in train
    csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
KeyError: 'model'

So what can I do to edit backbone net?

You need to rename the weights appropriately since the PyTorch model definition is different. You may need to add some extra wrapper to conform to what YOLOv5 train script expects. For example, it seems like the train script expects a checkpoint file, with a key model.

If you want to train a YOLOv5 object detector, I recommend you to use their pre-trained weights on COCO. Unless you have specific needs, using their COCO pre-trained weights will be better.

You need to rename the weights appropriately since the PyTorch model definition is different. You may need to add some extra wrapper to conform to what YOLOv5 train script expects. For example, it seems like the train script expects a checkpoint file, with a key model.

If you want to train a YOLOv5 object detector, I recommend you to use their pre-trained weights on COCO. Unless you have specific needs, using their COCO pre-trained weights will be better.

Thanks a lot, I have to get the dectecter without COCO pretrained.Maybe yolo don't suit currnet apply.
But thanks your comment!