加载预训练模型转onnx报错

[Tristesse-stk]

如题，报错如下：
Exception has occurred: RuntimeError
Error(s) in loading state_dict for YoloBody:
size mismatch for yolo_head_P3.weight: copying a param with shape torch.Size([21, 256, 1, 1]) from checkpoint, the shape in current model is torch.Size([258, 256, 1, 1]).
size mismatch for yolo_head_P3.bias: copying a param with shape torch.Size([21]) from checkpoint, the shape in current model is torch.Size([258]).
size mismatch for yolo_head_P4.weight: copying a param with shape torch.Size([21, 512, 1, 1]) from checkpoint, the shape in current model is torch.Size([258, 512, 1, 1]).
size mismatch for yolo_head_P4.bias: copying a param with shape torch.Size([21]) from checkpoint, the shape in current model is torch.Size([258]).
size mismatch for yolo_head_P5.weight: copying a param with shape torch.Size([21, 1024, 1, 1]) from checkpoint, the shape in current model is torch.Size([258, 1024, 1, 1]).
size mismatch for yolo_head_P5.bias: copying a param with shape torch.Size([21]) from checkpoint, the shape in current model is torch.Size([258]).

根据报错提示，作者本人提供基于SSDD数据集训练完成的pth类别数只有1类嘛？

[Egrt]

直接将ssdd_classes.txt中的类别修改为只有一类 ship就可以了。

[Tristesse-stk]

直接将ssdd_classes.txt中的类别修改为只有一类 ship就可以了。

好的，谢谢楼主，已解决。