NVIDIA-AI-IOT / yolo_deepstream

yolo model qat and deploy with deepstream&tensorrt
Apache License 2.0
533 stars 135 forks source link

_pickle.UnpicklingError: invalid load key, '<'. #22

Closed mfoglio closed 2 years ago

mfoglio commented 2 years ago

Within the docker container nvcr.io/nvidia/deepstream:6.0-triton:

mkdir /src
cd /src
git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
cd pytorch-YOLOv4
apt -y install python3-venv
python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip3 install \
numpy==1.18.2 \
torch==1.4.0 \
tensorboardX==2.0 \
scikit_image==0.16.2 \
matplotlib==2.2.3 \
tqdm==4.43.0 \
easydict==1.9 \
Pillow==7.1.2 \
opencv_python \
onnx \
onnxruntime

wget --no-check-certificate "https://docs.google.com/uc?export=download&id=1wv_LiFeCRYwtpkqREPeI13-gPELBDwuJ" -r -A 'uc*' -e robots=off -nd -O yolov4.pth
python3 demo_pytorch2onnx.py yolov4.pth data/dog.jpg 8 80 416 416

Error:

(venv) root@ip-172-31-9-127:/src/pytorch-YOLOv4# python3 demo_pytorch2onnx.py ./yolov4.pth data/dog.jpg 8 80 416 416
Converting to onnx and running demo ...
Traceback (most recent call last):
  File "demo_pytorch2onnx.py", line 96, in <module>
    main(weight_file, image_path, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W)
  File "demo_pytorch2onnx.py", line 72, in main
    transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W)
  File "demo_pytorch2onnx.py", line 19, in transform_to_onnx
    pretrained_dict = torch.load(weight_file, map_location=torch.device('cuda'))
  File "/src/pytorch-YOLOv4/venv/lib/python3.8/site-packages/torch/serialization.py", line 529, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/src/pytorch-YOLOv4/venv/lib/python3.8/site-packages/torch/serialization.py", line 692, in _legacy_load
    magic_number = pickle_module.load(f, **pickle_load_args)
_pickle.UnpicklingError: invalid load key, '<'.
mfoglio commented 2 years ago

Solved here https://github.com/Tianxiaomo/pytorch-YOLOv4/issues/495