Fix initializing anchor_grids in load_from_ultralytics

zhiqwang · Oct 20, 2021 · bdeb0f2 · bdeb0f2
1 parent 3ba4a2a
commit bdeb0f2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/yolort/utils/update_module_state.py b/yolort/utils/update_module_state.py
@@ -6,6 +6,7 @@
 
 from yolort.models import yolo
 from yolort.v5 import load_yolov5_model, get_yolov5_size
+from .image_utils import to_numpy
 
 
 def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"):
@@ -30,7 +31,9 @@ def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"):
     strides = checkpoint_yolov5.stride
     anchor_grids = checkpoint_yolov5.yaml["anchors"]
     if isinstance(anchor_grids, int):
-        anchor_grids = [list(range(anchor_grids * 2))] * len(strides)
+        anchor_grids = (
+            to_numpy(checkpoint_yolov5.model[-1].anchor_grid).reshape(3, -1).tolist()
+        )
 
     depth_multiple = checkpoint_yolov5.yaml["depth_multiple"]
     width_multiple = checkpoint_yolov5.yaml["width_multiple"]