将自己数据集训练好的pth文件转化为onnx文件报错
我输入的命令:python tools/export_onnx.py -n yolox-s -c weights/best_ckpt.pth --output-name weights/best_ckpt.onnx
报错内容:RuntimeError: Error(s) in loading state_dict for YOLOX:
size mismatch for head.cls_preds.0.weight: copying a param with shape torch.Size([6, 128, 1, 1]) from checkpoint, the shape in current model is torch.Size([80, 128, 1, 1]).
size mismatch for head.cls_preds.0.bias: copying a param with shape torch.Size([6]) from checkpoint, the shape in current model is torch.Size([80]).
size mismatch for head.cls_preds.1.weight: copying a param with shape torch.Size([6, 128, 1, 1]) from checkpoint, the shape in current model is torch.Size([80, 128, 1, 1]).
size mismatch for head.cls_preds.1.bias: copying a param with shape torch.Size([6]) from checkpoint, the shape in current model is torch.Size([80]).
size mismatch for head.cls_preds.2.weight: copying a param with shape torch.Size([6, 128, 1, 1]) from checkpoint, the shape in current model is torch.Size([80, 128, 1, 1]).
size mismatch for head.cls_preds.2.bias: copying a param with shape torch.Size([6]) from checkpoint, the shape in current model is torch.Size([80]).
我的解答思路:好像是我的数据集只有6种分类,但是export_onnx.py默认是80种分类。但是我不知道在哪改这个默认值
我想要达到的结果