search

bubbliiiing / yolov4-tiny-pytorch

这是一个YoloV4-tiny-pytorch的源码，可以用于训练自己的模型。
MIT License
787 stars 184 forks source link

使用log里面的训练生成的权重文件，出现以下错误 #72

Closed yifan0414 closed 2 years ago

yifan0414 commented 2 years ago

训练时用的是yolov4_tiny_weights_coco.pth权重文件

(myconda) root@c62f91247b39:/mnt/yolov4-tiny-pytorch# python predict.py
Traceback (most recent call last):
  File "predict.py", line 14, in <module>
    yolo = YOLO()
  File "/mnt/yolov4-tiny-pytorch/yolo.py", line 97, in __init__
    self.generate()
  File "/mnt/yolov4-tiny-pytorch/yolo.py", line 108, in generate
    self.net.load_state_dict(torch.load(self.model_path, map_location=device))
  File "/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for YoloBody:
        size mismatch for yolo_headP5.1.weight: copying a param with shape torch.Size([75, 512, 1, 1]) from checkpoint, the shape in current model is torch.Size([255, 512, 1, 1]).
        size mismatch for yolo_headP5.1.bias: copying a param with shape torch.Size([75]) from checkpoint, the shape in current model is torch.Size([255]).
        size mismatch for yolo_headP4.1.weight: copying a param with shape torch.Size([75, 256, 1, 1]) from checkpoint, the shape in current model is torch.Size([255, 256, 1, 1]).
        size mismatch for yolo_headP4.1.bias: copying a param with shape torch.Size([75]) from checkpoint, the shape in current model is torch.Size([255]).

yolo.py修改部分如下：

#--------------------------------------------------------------------------#
        #   使用自己训练好的模型进行预测一定要修改model_path和classes_path！
        #   model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
        #
        #   训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
        #   验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
        #   如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
        #--------------------------------------------------------------------------#
        # "model_path"        : 'model_data/yolov4_tiny_weights_coco.pth',
        "model_path"        : 'logs/newer.pth',
        "classes_path"      : 'model_data/coco_classes.txt',
image