训练模型正常，导入模型测试报错

[sjw-cool]

我用yolov4tiny训练自己的数据，总共也就2类。训练没有问题，可是在yolo.py中导入自己的训练模型就报torch.size不匹配的错，但是导入yolov4_tiny_weights_voc.pth和yolov4_tiny_weights_coco.pth就没有问题

[sjw-cool]

报错log如下:
ssh://root@10.31.15.32:22037/root/anaconda3/envs/pytorch/bin/python3 -u /root/PychramProject/yolov4-tiny-pytorch-master/predict.py
Loading weights into state dict...
Traceback (most recent call last):
File "/root/PychramProject/yolov4-tiny-pytorch-master/predict.py", line 7, in
yolo = YOLO()
File "/root/PychramProject/yolov4-tiny-pytorch-master/yolo.py", line 47, in init
self.generate()
File "/root/PychramProject/yolov4-tiny-pytorch-master/yolo.py", line 79, in generate
self.net.load_state_dict(state_dict)
File "/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1045, in load_state_dict
self.class.name, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for YoloBody:
size mismatch for yolo_headP5.1.weight: copying a param with shape torch.Size([21, 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([21]) 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([21, 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([21]) from checkpoint, the shape in current model is torch.Size([255]).

Process finished with exit code 1

[yifan0414]

遇到了同样的问题，请问为什么呢 @sjw-cool