Fix using_checkpoint in unicom download_fn; fix CLIP input_resolution

oml/models/vit_clip/extractor.py

@@ -54,7 +54,7 @@ def vitl14_224() -> VisionTransformer:
 def vitl14_336() -> VisionTransformer:
     return VisionTransformer(
         output_dim=768,
-        input_resolution=224,
+        input_resolution=336,
         layers=24,
         width=1024,
         patch_size=14,

oml/models/vit_unicom/external/model.py

@@ -83,7 +83,7 @@ def _download(url: str, root: str):
 
 
 # copy from https://github.com/openai/CLIP/blob/main/clip/clip.py#L94
-def load(name: str, device: str = "cpu", download_root: str = None):
+def load(name: str, device: str = "cpu", download_root: str = None, using_checkpoint: bool = True):
     if name in _MODELS:
         model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/unicom"))
     elif os.path.isfile(name):
@@ -93,7 +93,7 @@ def load(name: str, device: str = "cpu", download_root: str = None):
     with open(model_path, "rb") as opened_file:
         state_dict = torch.load(opened_file, map_location="cpu")
 
-    model, transform = load_model_and_transform(name)
+    model, transform = load_model_and_transform(name, using_checkpoint=using_checkpoint)
     state_dict_fp32 = {}
     for k, v in state_dict.items():
         state_dict_fp32[k] = v.float()

oml/models/vit_unicom/external/vision_transformer.py

@@ -347,14 +347,14 @@ def transform(im_size):
     )
 
 
-def load_model_and_transform(name="ViT-L/14@336px"):
+def load_model_and_transform(name="ViT-L/14@336px", using_checkpoint: bool = True):
     if name == "ViT-B/32":
-        return build_model(name), transform(224)
+        return build_model(name, using_checkpoint=using_checkpoint), transform(224)
     elif name == "ViT-B/16":
-        return build_model(name), transform(224)
+        return build_model(name, using_checkpoint=using_checkpoint), transform(224)
     elif name == "ViT-L/14":
-        return build_model(name), transform(224)
+        return build_model(name, using_checkpoint=using_checkpoint), transform(224)
     elif name == "ViT-L/14@336px":
-        return build_model(name), transform(336)
+        return build_model(name, using_checkpoint=using_checkpoint), transform(336)
     else:
         raise

oml/models/vit_unicom/extractor.py

@@ -40,25 +40,33 @@ class ViTUnicomExtractor(IExtractor):
 
     pretrained_models = {
         "vitb32_unicom": {
-            "download_fn": lambda: load("ViT-B/32", download_root=CKPT_SAVE_ROOT),
+            "download_fn": lambda using_checkpoint: load(
+                "ViT-B/32", download_root=CKPT_SAVE_ROOT, using_checkpoint=using_checkpoint
+            ),
             "init_args": {"arch": "vitb32_unicom", "normalise_features": True},
         },
         "vitb16_unicom": {
-            "download_fn": lambda: load("ViT-B/16", download_root=CKPT_SAVE_ROOT),
+            "download_fn": lambda using_checkpoint: load(
+                "ViT-B/16", download_root=CKPT_SAVE_ROOT, using_checkpoint=using_checkpoint
+            ),
             "init_args": {"arch": "vitb16_unicom", "normalise_features": True},
         },
         "vitl14_unicom": {
-            "download_fn": lambda: load("ViT-L/14", download_root=CKPT_SAVE_ROOT),
+            "download_fn": lambda using_checkpoint: load(
+                "ViT-L/14", download_root=CKPT_SAVE_ROOT, using_checkpoint=using_checkpoint
+            ),
             "init_args": {"arch": "vitl14_unicom", "normalise_features": True},
         },
         "vitl14_336px_unicom": {
-            "download_fn": lambda: load("ViT-L/14@336px", download_root=CKPT_SAVE_ROOT),
+            "download_fn": lambda using_checkpoint: load(
+                "ViT-L/14@336px", download_root=CKPT_SAVE_ROOT, using_checkpoint=using_checkpoint
+            ),
             "init_args": {"arch": "vitl14_336px_unicom", "normalise_features": True},
         },
     }
 
     def __init__(
-        self, weights: Optional[Union[Path, str]], arch: str, normalise_features: bool, use_gradiend_ckpt: bool = True
+        self, weights: Optional[Union[Path, str]], arch: str, normalise_features: bool, use_gradient_ckpt: bool = True
     ):
         """
         Args:
@@ -68,20 +76,22 @@ def __init__(
             arch: Might be one of ``vitb32_unicom``, ``vitb16_unicom``, ``vitl14_unicom``, ``vitl14_336px_unicom``.
              You can check all the available options in ``self.constructors``
             normalise_features: Set ``True`` to normalise output features
-            use_gradiend_ckpt: Whether to use gradient checkpointing inside VisionTransformer class.
+            use_gradient_ckpt: Whether to use gradient checkpointing inside VisionTransformer class.
         """
         assert arch in self.constructors
         super(IExtractor, self).__init__()
 
         self.arch = arch
         self.normalise_features = normalise_features
 
-        self.model = self.constructors[arch](using_checkpoint=use_gradiend_ckpt)
+        self.model = self.constructors[arch](using_checkpoint=use_gradient_ckpt)
 
         if weights is None:
             return
         elif weights in self.constructors:
-            self.model, _ = self.pretrained_models[weights]["download_fn"]()  # type: ignore
+            self.model, _ = self.pretrained_models[weights]["download_fn"](  # type: ignore
+                using_checkpoint=use_gradient_ckpt
+            )
         else:
             ckpt = torch.load(weights, map_location="cpu")
             state_dict = ckpt["state_dict"] if "state_dict" in ckpt else ckpt