Error loading custom yolov5-rt-stack model in C++ converted from yolov5_v4.0 - file not found: archive/constants.pkl #142
Comments
|
Going to try and test on some other models I've trained. I might try and re-train and edit some of the arguments that I pass during training (such as removing |
|
Update, looks like I was training off of yolov5 code as of March 14. yolov5-v4.0 was released on Jan 4. So I might have been in a weird hybrid between v4.0 and v5.0. I'm going to check out v4.0 of yolov5, retrain, and try again. |
|
Did you try to generate the from yolort.models import yolov5s
model_yolort = yolov5s()
model_yolort.model.load_state_dict(torch.load("yolov5-rt-stack-yolov5_v4-model.pt")) # Edited
# jit scripting the pytorch model
traced_model = torch.jit.script(model_yolort)
traced_model.save("yolov5-rt-stack-yolov5_v4-model.torchscript.pt")And then use the generated If the inference results on the pytorch backend between yolov5 v5.0 (or master branch) and yolort are consistent, that should be ok. Actually the model structure of v4.0 and v5.0 are very similar, in other words, the current version of yolort r4.0 is compatible with yolov5 v5.0 or master branch. I guess that yolov5 v5.0 fuse their model automatically, that's why we have problem to convert the yolov5 in the v5.0 and master branch. |
|
My "pipeline" pseudocode is as follows: best.pt = ultralytics-yolov5(dataset)
model = update_module_state_from_ultralytics(arch='yolov5s',
version='v4.0',
custom_path_or_model=/path/to/best.pt)
torch.save(model.state_dict(), /path/to/yolort-best.pt)
./yolo_inference --input_source /path/to/img --checkpoint /path/to/yolort-best.ptI believe you're saying I'm missing a step in my pipeline and it should be like ( best.pt = ultralytics-yolov5(dataset)
model = update_module_state_from_ultralytics(arch='yolov5s',
version='v4.0',
custom_path_or_model=/path/to/best.pt)
torch.save(model.state_dict(), /path/to/yolort-best.pt)
>>>from yolort.models import yolov5s
>>>updated_model = yolov5s()
>>>updated_model.load_state_dict(torch.load(/path/to/yolort-best.pt))
>>>model_script = torch.jit.script(updated_model)
>>>model_script.save(/path/to/yolort-best.torchscript.pt)
./yolo_inference --input_source /path/to/img --checkpoint /path/to/yolort-best.torchscript.ptIs that correct? |
Yep! And I guess there is one minor fix in your snippets, updated_model.model.load_state_dict(torch.load(/path/to/yolort-best.pt))Because we wrapped the vanilla converted yolov5 model in And we do the pre-processing (with |
|
Are you saying to go from updated_model.load_state_dict(torch.load(/path/to/yolort-best.pt))to updated_model.model.load_state_dict(torch.load(/path/to/yolort-best.pt)) ??? If so, I am getting the error |
Sorry @mattpopovich , Try the following instead updated_model.model.model.load_state_dict(torch.load(/path/to/yolort-best.pt)) The current model is wrapped one layer deeper than I thought ( |
|
|
|
Regardless, I don't think that is the issue. The conversion works successfully after using |
|
If you're interested in reproducing the above error, you can run the below You can get yolov5s-v4.0.pt and bus.jpg from those links. Once you get the file You should get the If you think the below script is handy, I'd be happy to clean it up and add it to the repo. It is mostly a duplicate of your Jupyter notebook though! Click to display convert_ultralytics_to_rt-stack.py
# When given a path to a yolov5-v4.0 weights file,
# will convert it to a yolov5-rt-stack weights file,
# check it for equivalence,
# then convert it to torchscript for inference via LibTorch in C++
# With inspiration from: https://github.com/zhiqwang/yolov5-rt-stack/blob/master/notebooks/inference-pytorch-export-libtorch.ipynb
# Author: Matt Popovich (mattpopovich.com)
import sys
import os
import torch
import cv2
sys.path.insert(0, "/path/to/yolov5-rt-stack")
from yolort.utils.image_utils import (
letterbox,
non_max_suppression,
plot_one_box,
scale_coords,
color_list,
)
from yolort.utils import (
cv2_imshow,
get_image_from_url,
read_image_to_tensor,
update_module_state_from_ultralytics,
)
# Define static variables
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0"
device = torch.device('cuda')
is_half = False
conf = 0.4
iou = 0.45
img_size = 512
# Define file paths
img_path = 'bus.jpg'
ultralytics_weights_path = 'yolov5s-v4.0.pt'
yolort_weight_path = 'yolov5s-v4.0-RT.pt'
yolort_script_path = 'yolov5s-v4.0-RT.torchscript.pt'
label_path = "coco.names"
# Check paths given
if not os.path.isfile(ultralytics_weights_path):
sys.exit("ERROR: filename specified does not exist: " + ultralytics_weights_path)
if os.path.isfile(img_path):
img_raw = cv2.imread(img_path)
else:
sys.exit("ERROR: filename specified for img_path does not exist: " + img_path)
with open(label_path) as f:
num_classes = 0
# Count number of lines that have text in them
for line in f:
# If line is empty, it will still contain '\n' = len of 1
num_classes += 1 if len(line) > 1 else 0
print("number of lines in {} = {}".format(label_path, num_classes))
# Preprocess
img = letterbox(img_raw, new_shape=(img_size,img_size))[0]
img = read_image_to_tensor(img, is_half)
img = img.to(device)
## Load model as ultralytics and inference
model = torch.hub.load('ultralytics/yolov5', 'custom', path=ultralytics_weights_path, autoshape=False)
model = model.to(device)
model.conf = conf # confidence threshold (0-1)
model.iou = iou # NMS IoU threshold (0-1)
model.classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for persons, cats and dogs
model.eval()
# Get actual anchors from ultralytics model
m = model.model[-1] # get Detect() layer
anchor_grids = m.anchor_grid.view((3, -1)).cpu().tolist() # get anchors
with torch.no_grad():
ultralytics_dets = model(img[None])[0]
ultralytics_dets = non_max_suppression(ultralytics_dets, conf, iou, agnostic=True)[0]
print(f'Detection results with ultralytics:\n{ultralytics_dets}')
model = update_module_state_from_ultralytics(arch='yolov5s',
version='v4.0',
custom_path_or_model=ultralytics_weights_path,
set_fp16=is_half,
num_classes=num_classes)
# The updated model checkpoint
torch.save(model.state_dict(), yolort_weight_path)
## Load model as yolort and inference
from yolort.models.yolo import yolov5_darknet_pan_s_r40 as yolov5s
model = yolov5s(score_thresh=conf, nms_thresh=iou, num_classes=num_classes, anchor_grids=anchor_grids)
model.load_state_dict(torch.load(yolort_weight_path))
model = model.to(device)
model.eval()
with torch.no_grad():
yolort_dets = model(img[None])
print(f"Detection boxes with yolort:\n{yolort_dets[0]['boxes']}")
print(f"Detection scores with yolort:\n{yolort_dets[0]['scores']}")
print(f"Detection labels with yolort:\n{yolort_dets[0]['labels']}")
# Testing boxes
torch.testing.assert_allclose(
yolort_dets[0]['boxes'], ultralytics_dets[:, :4], rtol=1e-05, atol=1e-07)
# Testing scores
torch.testing.assert_allclose(
yolort_dets[0]['scores'], ultralytics_dets[:, 4], rtol=1e-05, atol=1e-07)
# Testing labels
torch.testing.assert_allclose(
yolort_dets[0]['labels'], ultralytics_dets[:, 5].to(dtype=torch.int64), rtol=1e-05, atol=1e-07)
print("Exported model has been tested, and the result looks good!")
## Detect output visualization
# Get label names
from yolort.utils.image_utils import load_names
LABELS = load_names(label_path)
COLORS = color_list()
# Hah, that's the trick to rescale the box correctly
boxes = scale_coords(yolort_dets[0]['boxes'], img.shape[1:], img_raw.shape[:-1])
for box, label in zip(boxes.tolist(), yolort_dets[0]['labels'].tolist()):
img_raw = plot_one_box(box, img_raw, color=COLORS[label % len(COLORS)], label=LABELS[label])
# cv2_imshow(img_raw, imshow_scale=0.5) # If in Jupyter notebook
# cv2.imshow("img_raw", img_raw) # If running script locally
# cv2.waitKey()
# Create script of yolov5-rt-stack model
print("\nBeginning export of torchscript model...")
model_script = torch.jit.script(model)
model_script.eval()
model_script = model_script.to(device)
model_script.save(yolort_script_path)
print("Exported yolov5-rt-stack torchscript model to: ")
print("\t" + yolort_script_path) |
|
Hi @mattpopovich , Thank you for your feedback on this phenomenon, I'll try to reproduce this bug and try to fix this problem ASAP. One thing that can be determined with a high probability is that you don't need to check out v4.0 of yolov5, retrain your model . |
|
Hi, @mattpopovich , Thanks for your details analysis and scripts above, I think I got where our operations are different, there is one thing we need to be careful when calling from yolort.models.yolo import yolov5_darknet_pan_s_r40 as yolov5s # This is `YOLO`
model = yolov5s(score_thresh=conf, nms_thresh=iou, num_classes=num_classes, anchor_grids=anchor_grids)
model.load_state_dict(torch.load(yolort_weight_path))
model = model.to(device)These lines should be changed to from yolort.models import yolov5s # This is `YOLOModule`
# Make sure you are using the master branch here!
model_yolort = yolov5s(
pretrained=False,
num_classes=num_classes,
anchor_grids=anchor_grids, # same as above
nms_thresh=0.45,
score_thresh=0.3,
)
# Load your updated checkpoint here
model_yolort.model.load_state_dict(torch.load("yolov5s-v4.0-RT.pt"))
# jit scripting the pytorch model
model_yolort_script = torch.jit.script(model_yolort)
model_yolort_script.save("yolov5s-v4.0-RT.torchscript.pt")Let me know if the above modification works for you! |
|
FYI, The pre-processing operation is in There are some minor differences in the behavior of This difference should be able to be eliminated, but we do not have time to resolve this problem now. That's why we only test the behavior between And you can see #92 (comment) for the difference in mAP between BTW, the post-processing ( All contributions and suggestions are welcome here. |
Great news, I ran a quick test tonight and it appears to work on the stock yolov5-v4.0 pre-trained MS COCO weights file! I will do further validation tomorrow.
That was definitely the issue. I remember not remembering the path of how to import yolov5s and just grabbed a line from the first Jupyter notebook that I had open. I did not realize that they ( Really appreciate your extremely speedy responses in getting this resolved! Awesome repo you have here. |
🐛 Bug
After training a yolov5-v4.0 model, I then took its
best.ptweights and converted them to yolov5-rt-stack weights via python inupdate_module_state_from_ultralytics(). I then took these yolov5-rt-stack weights and passed them as an argument to the stock./yolo_inferenceprogram.Is that the correct procedure?
I get the following error:
I do not have any errors and everything runs correctly when using
yolov5-rt-stack/test/tracing/yolov5s.torchscript.ptweights with./yolo_inference, so I believe I have everything installed, compiling, and running correctly.Expected behavior
./yolo_inferenceruns correctly, finishes, and outputs detections.Environment
Click to display environment
conda,pip, source): source / sourceAdditional context
I trained my yolov5_v4.0 model with the following command (on 4x Tesla V100's):
train.py --local_rank=0 --img 512 --batch 128 --epochs 500 --data path/to/yaml --weights --cfg models/yolov5s.yaml --device 0,1,2,3 --name experiment-2The text was updated successfully, but these errors were encountered: