Tianxiaomo / pytorch-YOLOv4

PyTorch ,ONNX and TensorRT implementation of YOLOv4
Apache License 2.0
4.46k stars 1.49k forks source link

Update models.py #546

Open Rookiecoder-Wang opened 1 year ago

Rookiecoder-Wang commented 1 year ago

A Very Easy Bug Fix(A Rookie's First Test) For This Error RuntimeError: Error(s) in loading state_dict for Yolov4:

Rookiecoder-Wang commented 1 year ago

If without this fix,an error will occured when we running demo_pytorch2onnx.py,just like this Traceback (most recent call last): File "demo_pytorch2onnx.py", line 96, in main(weight_file, image_path, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W) File "demo_pytorch2onnx.py", line 69, in main onnx_path_demo = transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W) File "demo_pytorch2onnx.py", line 20, in transform_to_onnx model.load_state_dict(pretrained_dict) File "F:\Pythons\envs\modelt\lib\site-packages\torch\nn\modules\module.py", line 830, in load_state_dict self.class.name, "\n\t".join(error_msgs))) RuntimeError: Error(s) in loading state_dict for Yolov4: Missing key(s) in state_dict: "neck.conv1.conv.0.weight", "neck.conv1.conv.1.weight", "neck.conv1.conv.1.bias", "neck.conv1.conv.1.running_mean", "neck.conv1.conv.1.running_var", "neck.conv2.conv.0.weight", "neck.conv2.conv.1.weight", "neck.conv2.conv.1.bias", "neck.conv2.conv.1.running_mean", "neck.conv2.conv.1.running_var", "neck.conv3.conv.0.weight", "neck.conv3.conv.1.weight", "neck.conv3.conv.1.bias", "neck.conv3.conv.1.running_mean", "neck.conv3.conv.1.running_var", "neck.conv4.conv.0.weight", "neck.conv4.conv.1.weight", "neck.conv4.conv.1.bias", "neck.conv4.conv.1.running_mean", "neck.conv4.conv.1.running_var", "neck.conv5.conv.0.weight", "neck.conv5.conv.1.weight", ········