Model hot fix

Model hot fix

Author: AlekseySh

oml/models/utils.py

@@ -13,7 +13,12 @@ def remove_criterion_in_state_dict(state_dict: TStateDict) -> TStateDict:
 
 
 def find_prefix_in_state_dict(state_dict: TStateDict, trial_key: str) -> str:
-    k0 = [k for k in state_dict.keys() if trial_key in k][0]
+    keys_starting_with_trial_key = [k for k in state_dict.keys() if trial_key in k]
+    assert keys_starting_with_trial_key, (
+        f"There are no keys starting from {trial_key}.\n" f"The existing keys are: {list(state_dict.keys())}"
+    )
+
+    k0 = keys_starting_with_trial_key[0]
     prefix = k0[: k0.index(trial_key)]
 
     keys_not_starting_with_prefix = list(filter(lambda x: not x.startswith(prefix), state_dict.keys()))

oml/models/vit_clip/extractor.py

@@ -151,13 +151,13 @@ def __init__(
             visual = torch.jit.load(Path(weights), map_location="cpu").visual
             patch_device_and_float(visual, device="cpu")
             state_dict = visual.state_dict()
+
         else:
             state_dict = torch.load(Path(weights), map_location="cpu")
             state_dict = state_dict.get("state_dict", state_dict)
-
-        state_dict = remove_criterion_in_state_dict(state_dict)
-        state_dict = remove_prefix_from_state_dict(state_dict, trial_key="conv1.weight")
-        state_dict = take_visual_part_of_vit_clip(state_dict, needed_keys=self.visual.state_dict().keys())
+            state_dict = remove_criterion_in_state_dict(state_dict)
+            state_dict = take_visual_part_of_vit_clip(state_dict, needed_keys=self.visual.state_dict().keys())
+            state_dict = remove_prefix_from_state_dict(state_dict, trial_key="conv1.weight")
 
         self.visual.load_state_dict(state_dict=state_dict, strict=True)
 
@@ -178,8 +178,9 @@ def feat_dim(self) -> int:
 
 def take_visual_part_of_vit_clip(state_dict: TStateDict, needed_keys: Iterable[str]) -> TStateDict:
     for k in list(state_dict):
-        if k.startswith("visual."):
-            state_dict[k.lstrip("visual")[1:]] = state_dict.pop(k)
+        if "visual" in k:
+            new_key = k[k.find("visual") + len("visual") + 1 :]
+            state_dict[new_key] = state_dict.pop(k)
     state_dict = filter_state_dict(state_dict, needed_keys=needed_keys)
     return state_dict