From 49cb27ccf886b17acfab66e53ea5d05b58daa6fe Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Wed, 16 Jun 2021 14:15:24 -0400 Subject: [PATCH] Update README about alignment with ultralytics --- README.md | 2 +- yolort/utils/update_module_state.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index c360fe3e..afaffbcf 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ model = torch.hub.load('zhiqwang/yolov5-rt-stack', 'yolov5s', pretrained=True) ### Updating checkpoint from ultralytics/yolov5 -The module state of `yolort` has some differences comparing to `ultralytics/yolov5`. We can load ultralytics's trained model checkpoint with minor changes, and we have converted ultralytics's release [v3.1](https://github.com/ultralytics/yolov5/releases/tag/v3.1) and [v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0). For example, if you want to convert a `yolov5s` (release 4.0) model, you can just run the following script: +The module state of `yolort` has some differences comparing to `ultralytics/yolov5`. We can load ultralytics's trained model checkpoint with minor changes, and we have converted ultralytics's release [v3.1](https://github.com/ultralytics/yolov5/releases/tag/v3.1) and [v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0). For example, if you want to convert a `yolov5s` (release 4.0) model, you can just run the following script. You can also see our [how-to-align-with-ultralytics-yolov5](http://github.com/zhiqwang/yolov5-rt-stack/blob/master/notebooks/how-to-align-with-ultralytics-yolov5.ipynb) notebook for more details. ```python from yolort.utils import update_module_state_from_ultralytics diff --git a/yolort/utils/update_module_state.py b/yolort/utils/update_module_state.py index fbf6cb7e..a4da922b 100644 --- a/yolort/utils/update_module_state.py +++ b/yolort/utils/update_module_state.py @@ -61,8 +61,9 @@ def update_module_state_from_ultralytics( model = torch.hub.load(f'ultralytics/yolov5:{version}', 'custom', path_or_model=custom_path_or_model) key_arch = f'{arch}_{feature_fusion_type.lower()}_{version}' - assert key_arch in architecture_maps, ( - "Currently does't support your architecture, fell free to report a issue to yolort") + if key_arch not in architecture_maps: + raise ValueError("Currently does't supports this architecture, " + "fell free to file an issue labeled enhancement to us") module_state_updater = ModuleStateUpdate( arch=architecture_maps[key_arch],