🔀 [Merge] remote-tracking branch 'origin/fix/clean-load-weight-log'

Author: henrytsui000

yolo/config/task/train.yaml

@@ -51,4 +51,4 @@ scheduler:
 
 ema:
   enable: true
-  decay: 0.995
+  decay: 0.9999

yolo/model/yolo.py

@@ -133,9 +133,8 @@ def save_load_weights(self, weights: Union[Path, OrderedDict]):
         """
         if isinstance(weights, Path):
             weights = torch.load(weights, map_location=torch.device("cpu"), weights_only=False)
-        if "model_state_dict" in weights:
-            weights = weights["model_state_dict"]
-
+        if "state_dict" in weights:
+            weights = {name.removeprefix("model.model."): key for name, key in weights["state_dict"].items()}
         model_state_dict = self.model.state_dict()
 
         # TODO1: autoload old version weight
@@ -152,8 +151,15 @@ def save_load_weights(self, weights: Union[Path, OrderedDict]):
             model_state_dict[model_key] = weights[model_key]
 
         for error_name, error_set in error_dict.items():
-            for weight_name in error_set:
-                logger.warning(f":warning: Weight {error_name} for key: {'.'.join(weight_name)}")
+            error_dict = dict()
+            for layer_idx, *layer_name in error_set:
+                if layer_idx not in error_dict:
+                    error_dict[layer_idx] = [".".join(layer_name)]
+                else:
+                    error_dict[layer_idx].append(".".join(layer_name))
+            for layer_idx, layer_name in error_dict.items():
+                layer_name.sort()
+                logger.warning(f":warning: Weight {error_name} for Layer {layer_idx}: {', '.join(layer_name)}")
 
         self.model.load_state_dict(model_state_dict)
 

yolo/tools/format_converters.py

@@ -1,35 +1,36 @@
+convert_dict = {
+    "19.cv1": "19.conv",
+    "16.cv1": "16.conv",
+    ".7.cv1": ".7.conv",
+    ".5.cv1": ".5.conv",
+    ".3.cv1": ".3.conv",
+    ".28.": ".29.",
+    ".25.": ".26.",
+    ".22.": ".23.",
+    "cv": "conv",
+    ".m.": ".bottleneck.",
+}
+
+HEAD_NUM = "29"
+
+
 def convert_weight(old_state_dict, new_state_dict, model_size: int = 38):
-    # TODO: need to refactor
-    shift = 1
-    for idx in range(model_size):
-        new_list, old_list = [], []
-        for weight_name, weight_value in new_state_dict.items():
-            if weight_name.split(".")[0] == str(idx):
-                new_list.append((weight_name, None))
-        for weight_name, weight_value in old_state_dict.items():
-            if f"model.{idx+shift}." in weight_name:
-                old_list.append((weight_name, weight_value))
-        if len(new_list) == len(old_list):
-            for (weight_name, _), (_, weight_value) in zip(new_list, old_list):
-                new_state_dict[weight_name] = weight_value
+    new_weight_set = set(new_state_dict.keys())
+    for weight_name, weight_value in old_state_dict.items():
+        if HEAD_NUM in weight_name:
+            _, _, conv_name, conv_id, *post_fix = weight_name.split(".")
+            head_id = 30 if conv_name in ["cv2", "cv3"] else 22
+            head_type = "anchor_conv" if conv_name in ["cv2", "cv4"] else "class_conv"
+            weight_name = ".".join(["model", str(head_id), "heads", conv_id, head_type, *post_fix])
         else:
-            for weight_name, weight_value in old_list:
-                if "dfl" in weight_name:
-                    continue
-                _, _, conv_name, conv_idx, *details = weight_name.split(".")
-                if conv_name == "cv4" or conv_name == "cv5":
-                    layer_idx = 22
-                    shift = 2
-                else:
-                    layer_idx = 37
-
-                if conv_name == "cv2" or conv_name == "cv4":
-                    conv_task = "anchor_conv"
-                if conv_name == "cv3" or conv_name == "cv5":
-                    conv_task = "class_conv"
-
-                weight_name = ".".join([str(layer_idx), "heads", conv_idx, conv_task, *details])
-                new_state_dict[weight_name] = weight_value
+            for old_name, new_name in convert_dict.items():
+                if old_name in weight_name:
+                    weight_name = weight_name.replace(old_name, new_name)
+        if weight_name in new_weight_set:
+            assert new_state_dict[weight_name].shape == weight_value.shape, "shape miss match"
+            new_state_dict[weight_name] = weight_value
+            new_weight_set.remove(weight_name)
+
     return new_state_dict