No module named 'yolort.utils.update_module_state' while saving the Yolo Model

[HamzaAsiff]

Hello,
I followed the conversation and code snippet shared in "Error when convert custom model #90" but when I try to run the code I get the following error:

No module named 'yolort.utils.update_module_state'

Code Snippet:

from yolort.utils.update_module_state import ModuleStateUpdate

model = torch.hub.load(
    'ultralytics/yolov5',
    'custom',
    path_or_model='best.pt',  # I download this from Google drive provided by you.
)

module_state_updater = ModuleStateUpdate(arch='yolov5_darknet_pan_s_r40', num_classes=1)
module_state_updater.updating(model)
module_updated = module_state_updater.model

torch.save(module_updated.state_dict(), 'yolov5s_updated.pt')

I installed the latest version of yolort and check the repo but I couldn't find the update_modue_state.py function in the utils directory. Was the file in some previous yolort release?

---

Code format update by Zhiqiang

[zhiqwang]

Hi @HamzaAsiff , This code snippets is now deprecated, and we provide a command-line tool to export yolort-style model files instead

python tools/convert_yolov5_to_yolort.py --checkpoint_path {path/to/yolov5s.pt}

[HamzaAsiff]

@zhiqwang thanks for the quick reply
Can you please tell where can I find this "tools/convert_yolov5_to_yolort.py"?

[zhiqwang]

Here https://github.com/zhiqwang/yolov5-rt-stack/blob/main/tools/convert_yolov5_to_yolort.py

[HamzaAsiff]

I'm getting the following error:

$ python tools/convert_yolov5_to_yolort.py --checkpoint_path yoloBest.pt
Command Line Args: Namespace(checkpoint_path='yoloBest.pt', image_path='zidane.jpg', output_path=None, version='r6.0')
Traceback (most recent call last):
  File "tools/convert_yolov5_to_yolort.py", line 42, in <module>
    cli_main()
  File "tools/convert_yolov5_to_yolort.py", line 38, in cli_main
    convert_yolov5_checkpoint(checkpoint_path, output_path, version=args.version)
  File "C:\Users\Hamza\anaconda3\envs\FL2\lib\site-packages\yolort\models\_checkpoint.py", line 116, in convert_yolov5_checkpoint
    model_info = load_from_ultralytics(checkpoint_path, version=version)
  File "C:\Users\Hamza\anaconda3\envs\FL2\lib\site-packages\yolort\models\_checkpoint.py", line 33, in load_from_ultralytics
    checkpoint_yolov5 = load_yolov5_model(checkpoint_path)
  File "C:\Users\Hamza\anaconda3\envs\FL2\lib\site-packages\yolort\v5\helper.py", line 71, in load_yolov5_model
    model = ckpt["ema" if ckpt.get("ema") else "model"].float().eval()
KeyError: 'model'

Now I know why I'm getting this error. It's because the program is expecting the pt file in collections.OrderedDict format where model is a key in the dictionary. But what I'm passing is like this:

OrderedDict([('model.0.conv.weight',
              tensor([[[[ 1.30654e-03,  4.56353e-02,  4.78757e-02,  1.03541e-01,  5.85952e-02,  5.11846e-02],
                        [ 3.40025e-02,  8.85132e-02,  1.00922e-01,  1.36768e-01,  8.84443e-02,  6.07322e-02],
                        [ 2.17492e-02,  8.52733e-02,  1.34149e-01,  1.68065e-01,  7.17619e-02,  2.50236e-02],
                        [ 2.83325e-02,  6.81083e-02,  1.04575e-01,  1.24980e-01,  2.08013e-02, -1.42352e-02],
                        [-5.29770e-02, -2.78327e-02, -1.52003e-02, -5.02368e-03, -8.73413e-02, -1.14019e-01],
                        [-9.86468e-02, -9.80953e-02, -1.19396e-01, -8.50665e-02, -1.29254e-01, -1.22292e-01]]

What I want to achieve is the following:

1. Use the pretrained yolov5s.pt and finetune the model on coco dataset.
2. Save the weight for the model.
3. Save weight model as .pt file
4. Load the saved .pt file from step 4 and pass that to yolo train function for fine tuning.
5. Repeat step 1-4 until specified number of rounds is reached.

I'm getting the above error in the step 4.

[zhiqwang]

Now I know why I'm getting this error. It's because the program is expecting the pt file in collections.OrderedDict format where model is a key in the dictionary.

Yep, It seems to crash from following line, and you can customize this part of the code

https://github.com/zhiqwang/yolov5-rt-stack/blob/0c87bda491eb9bef2d9ee7a5f38112196ac80be1/yolort/models/_checkpoint.py#L116

[HamzaAsiff]

I managed to make a fix for this. The yolo train function expects a dictionary which includes:
{'epoch', 'best_fitness', 'model', 'ema', 'updates', 'optimizer', 'wandb_id', 'date'} but I only have the model part.
I passed the trained model weights in the .load_state_dict function to get the model and passed that into the model key value in the dictionary and set all of the other key values except model and date as None because even in the small and large variants of yolov5 these are none and passed that dictionary to the train function and it worked.

Thanks for the help, closing this issue now.