将自己数据集训练好的pth文件转化为onnx文件报错

将自己数据集训练好的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种分类。但是我不知道在哪改这个默认值

我想要达到的结果