Some fixes for CLIP

Some fixes for CLIP

Author: AlekseySh

oml/models/utils.py

@@ -16,7 +16,10 @@ 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]
     prefix = k0[: k0.index(trial_key)]
 
-    assert all(k.startswith(prefix) for k in state_dict.keys())
+    keys_not_starting_with_prefix = list(filter(lambda x: not x.startswith(prefix), state_dict.keys()))
+    assert (
+        not keys_not_starting_with_prefix
+    ), f"There are keys not starting from the found prefix {prefix}: {keys_not_starting_with_prefix}"
 
     return prefix
 

oml/models/vit_clip/extractor.py

@@ -93,13 +93,6 @@ class ViTCLIPExtractor(IExtractor):
             "fname": "openai_vitl14_224.ckpt",
             "init_args": {"arch": "vitl14_224", "normalise_features": False},
         },
-        "openai_vitl14_336": {
-            "url": f"{_OPENAI_URL}/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
-            "hash": "b311058cae50cb10fbfa2a44231c9473",
-            "is_jitted": True,
-            "fname": "openai_vitl14_336.ckpt",
-            "init_args": {"arch": "vitl14_336", "normalise_features": False},
-        },
         # checkpoints pretrained by SberbankAI
         "sber_vitb16_224": {
             "url": f"{_SBER_URL}/ruclip-vit-base-patch16-224/resolve/main/pytorch_model.bin",
@@ -161,9 +154,10 @@ def __init__(
         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="class_embedding")
-            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 = 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())
 
         self.visual.load_state_dict(state_dict=state_dict, strict=True)
 

oml/registry/transforms.py

@@ -67,7 +67,6 @@ def get_transforms_by_cfg(cfg: TCfg) -> TTransforms:
     "openai_vitb32_224": get_normalisation_resize_albu_clip(im_size=224),
     "openai_vitb16_224": get_normalisation_resize_albu_clip(im_size=224),
     "openai_vitl14_224": get_normalisation_resize_albu_clip(im_size=224),
-    "openai_vitl14_336": get_normalisation_resize_albu_clip(im_size=224),
     "vits16_inshop": get_normalisation_resize_hypvit(im_size=224, crop_size=224),
     "vits16_sop": get_normalisation_resize_hypvit(im_size=224, crop_size=224),
     "vits16_cars": get_normalisation_resize_albu(im_size=224),