yolox_nano.pth转yolox_nano.onnx文件报错:Missing key(s) & Unexpected key(s) in state_dict..(自己的数据集训练的模型)

我最近使用自己的数据集(只有1类)训练好了yolox_nano.pth模型,想把这个yolox_nano.pth转换成yolox_nano.onnx,但是报了很多错误:
转换命令(在YOLOX目录下):

python tools/export_onnx.py -n yolox-nano -c weights/yolox_nano.pth --output-name weights/yolox_nano.onnx

报错信息:

RuntimeError: Error(s) in loading state_dict for DataParallel:
 Missing key(s) in state_dict: ....
Unexpected key(s) in state_dict...

报错图如下:

img

img

我认为报错的原因是因为我转换的是自己的数据集,而我转换coco数据集的时候却不会报错,可以正常转换:

img

img

但我已经修改了yolox_base.py里面的num_classes、depth、width还有input_size这些:

img

还是报这个错误

完整报错信息如下:

(torch_G) E:\YOLOX>python tools/export_onnx.py -n yolox-nano -c weights/yolox_nano.pth --output-name weights/yolox_nano.onnx
2021-10-24 17:15:37.712 | INFO     | __main__:main:59 - args value: Namespace(batch_size=1, ckpt='weights/yolox_nano.pth', dynamic=False, exp_file=None, experiment_name=None, input='images', name='yolox-nano', no_onnxsim=False, opset=11, opts=[], output='output', output_name='weights/yolox_nano.onnx')
2021-10-24 17:15:40.278 | ERROR    | __main__:<module>:119 - An error has been caught in function '<module>', process 'MainProcess' (66764), thread 'MainThread' (59796):
Traceback (most recent call last):

> File "tools\export_onnx.py", line 119, in <module>
    main()
    └ <function main at 0x00000276616C75E0>

  File "tools\export_onnx.py", line 81, in main
    model.load_state_dict(ckpt)
    │     │               └ OrderedDict([('backbone.backbone.stem.conv.conv.weight', tensor([[[[ 1.2789e-02,  1.7050e-02,  2.4743e-02],
    │     │                           [ 6.535...
    │     └ <function Module.load_state_dict at 0x00000276609BEAF0>
    └ DataParallel(
        (module): YOLOX(
          (backbone): YOLOPAFPN(
            (backbone): CSPDarknet(
              (stem): Focus(
                (c...

  File "E:\Anaconda3\envs\torch_G\lib\site-packages\torch\nn\modules\module.py", line 1223, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(

RuntimeError: Error(s) in loading state_dict for DataParallel:
        Missing key(s) in state_dict: "module.backbone.backbone.stem.conv.conv.weight", "module.backbone.backbone.stem.conv.bn.weight", "module.backbone.backbone.stem.conv.bn.bias"......
        Unexpected key(s) in state_dict: "backbone.backbone.stem.conv.conv.weight", "backbone.backbone.stem.conv.bn.weight", "backbone.backbone.stem.conv.bn.bias"....

我认为是网络的问题,多了或者少了什么参数,导致不匹配。
我训练好的的yolox_nano.pth是没问题的,可以正常预测(推理):

img

img


那就感觉很奇怪了,希望有可以有这方面经验的同学可以帮我回答一下,谢谢。

看报错是你参数没修改?类别数目不对吧,原来的coco数据集是80类,你就只有一类,改下类别参数数目看看,不然不会说你转原始的没问题,转你训练的出问题了,这两个的区别就是类别数目不一样,然后就是一些参数不一样,但是网络结构之类都是一致的。

img

最后的报错看起来应该是原来的模型是分布式训练的(DDP),后面你是单卡,导致少了model.
如果采用修改pth的方案,给你写个简单的修改

    ckpt = torch.load(args.checkpoint, map_location=’cpu‘) #将保存的权重映射到cpu上
    ######################## here ############################
    from collections import OrderedDict
    sd = OrderedDict()
    for item in ckpt.items():
        sd_key = "model."+item[0] #将model.加在原来网络层名字之前
        sd[sd_key] = item[1]
    model.load_state_dict(sd,strict=True)

我的观点,你训练好的模型,如果在缓存里直接保存或者转换应该是正常的,但是如果你重新载入,然后再保存或者转换,会有问题的。
主要是有些参数模型并不知道,你需要先载入跟原来训练时一样的参数,然后在加载模型,比如比如学习率,激活函数等等,加载这些参数后,不需要训练,然后将训练好的模型加载,然后转换即可。

如果有帮助,麻烦采纳,谢谢!