TinyCLIP integration for ActionCLIP

clip/clip.py

@@ -7,7 +7,6 @@
 import urllib
 import warnings
 from typing import Union, List
-
 import torch
 from PIL import Image
 from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
@@ -20,22 +19,35 @@
 _tokenizer = _Tokenizer()
 
 _MODELS = {
-    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
-    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt"
-}
-
-def _download(url: str, root: str = os.path.expanduser("~/.cache/clip")):
+    "CLIP-ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "CLIP-ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "TinyCLIP-ViT-40M/32-Text-19M": "https://github.com/wkcn/TinyCLIP-model-zoo/releases/download/checkpoints/TinyCLIP-ViT-40M-32-Text-19M-LAION400M.pt",
+    "TinyCLIP-ViT-61M/32-Text-29M": "https://github.com/wkcn/TinyCLIP-model-zoo/releases/download/checkpoints/TinyCLIP-ViT-61M-32-Text-29M-LAION400M.pt",
+}   
+
+def _download(url: str, root: str = os.path.expanduser("~/.cache/clip"), skip_sha256 = False):
+    '''Download CLIP model to cache. Optionally, skip_sha256 checksum skips
+    the checksum to allow for downloading the official TinyCLIP models because 
+    they are not supplied.'''
+
+    if(skip_sha256):
+        warnings.warn(f"Skipping sha256 checksum matching because it is not available for TinyCLIP model.")
+
     os.makedirs(root, exist_ok=True)
     filename = os.path.basename(url)
 
     expected_sha256 = url.split("/")[-2]
     download_target = os.path.join(root, filename)
-
+        
     if os.path.exists(download_target) and not os.path.isfile(download_target):
         raise RuntimeError(f"{download_target} exists and is not a regular file")
-
+    
     if os.path.isfile(download_target):
-        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+        # TinyCLIP does not come with checksum, return
+        # download target early
+        if skip_sha256: 
+            return download_target
+        elif hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
             return download_target
         else:
             warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
@@ -46,12 +58,11 @@ def _download(url: str, root: str = os.path.expanduser("~/.cache/clip")):
                 buffer = source.read(8192)
                 if not buffer:
                     break
-
                 output.write(buffer)
                 loop.update(len(buffer))
-
-    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
-        raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+    if not skip_sha256:
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+            raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
 
     return download_target
 
@@ -93,8 +104,15 @@ def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_a
     preprocess : Callable[[PIL.Image], torch.Tensor]
         A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
     """
+    # Obtain CLIP backbone name by splitting on '-' 
+    # and obtaining the preprended CLIP name
+    clip_name = str.split(name, "-")[0] # [CLIP, TinyCLIP]
+
     if name in _MODELS:
-        model_path = _download(_MODELS[name])
+        is_tinyclip = clip_name == "TinyCLIP"
+        # Only skip sha_256 checksum if backbone set to TinyCLIP
+        # because no sha_256 checksum is available
+        model_path = _download(_MODELS[name], skip_sha256=is_tinyclip)
     elif os.path.isfile(name):
         model_path = name
     else:
@@ -112,8 +130,8 @@ def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_a
         state_dict = torch.load(model_path, map_location="cpu")
 
     if not jit:
-
-        model = build_model(state_dict or model.state_dict(), joint=joint,tsm=tsm,T=T,dropout=dropout, emb_dropout=emb_dropout,pretrain=pretrain).to(device)
+        # Split on first '-' to decide whether CLIP or TinyCLIP is backbone
+        model = build_model(state_dict or model.state_dict(), joint=joint,tsm=tsm,T=T,dropout=dropout, emb_dropout=emb_dropout,pretrain=pretrain, clip_backbone=clip_name).to(device)
         if str(device) == "cpu":
             model.float()
 

clip/model.py

@@ -296,12 +296,34 @@ def _convert_weights_to_fp16(l):
                     attr.data = attr.data.half()
 
     model.apply(_convert_weights_to_fp16)
-
-
-def build_model(state_dict: dict, tsm=False,T=8,dropout=0., joint=False,emb_dropout=0.,pretrain=True):
-    vit = "visual.proj" in state_dict
-
-    if vit:
+
+def convert_statedict(dict: dict, ResNet: bool = False) -> dict:
+    """Convert each key (component) in the TinyViT state_dict
+    to match each key (component) in the vanilla CLIP model
+    For example: 'module.visual.conv0.weight' -> 'visual.conv1.weight'"""    
+
+    new_dict = dict['state_dict'] # TinyCLIPs dict is ravelled 
+    if(ResNet):  
+        raise NotImplementedError("Conversion from TinyCLIP ResNet statedict not implemented, use a TinyClip ViT-based model instead.")
+    else:    
+        state_dict = {key.replace('module.', '', 1): value for key, value in new_dict.items()} # remove the 'module.' prefix of every key
+    return state_dict
+
+def build_model(state_dict: dict, tsm=False,T=8,dropout=0., joint=False,emb_dropout=0.,pretrain=True, clip_backbone = "CLIP"):      
+
+    # Determine selected clip backbone, convert
+    # state_dict weights of TinyCLIP to CLIP if
+    # TinyCLIP is selected
+    if(clip_backbone == "TinyCLIP"):
+        state_dict = convert_statedict(state_dict)
+    elif(clip_backbone == "CLIP"):
+        pass
+    else:
+        raise RuntimeError("build_model(): did not recognise CLIP backbone: {}, ensure model names are prepended with 'CLIP' or 'TinyCLIP'".format(clip_backbone))
+
+    print("Clip model (Vanilla clip/TinyCLIP) set to: {}".format(clip_backbone))
+    is_vit = 'visual.proj' in state_dict 
+    if is_vit:
         vision_width = state_dict["visual.conv1.weight"].shape[0]
         vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
         vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
@@ -330,7 +352,7 @@ def build_model(state_dict: dict, tsm=False,T=8,dropout=0., joint=False,emb_drop
         context_length, vocab_size, transformer_width, transformer_heads, transformer_layers,  tsm=tsm,T=T,joint=joint,
         dropout=dropout, emb_dropout=emb_dropout
     )
-
+    
     for key in ["input_resolution", "context_length", "vocab_size"]:
         if key in state_dict:
             del state_dict[key]
@@ -364,5 +386,5 @@ def build_model(state_dict: dict, tsm=False,T=8,dropout=0., joint=False,emb_drop
                 state_dict.pop(k)
 
         model.load_state_dict(state_dict,strict=False)
-
-    return model.eval()
+        
+    return model.eval()

configs/hmdb51/hmdb_test.yaml

@@ -18,7 +18,7 @@ data:
         N: 0 #2
         M: 0  #9
 network:
-    arch: ViT-B/16  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True 
     drop_out: 0.0 
     emb_dropout: 0.0

configs/hmdb51/hmdb_train.yaml

@@ -21,7 +21,7 @@ data:
         N: 0 #2
         M: 0  #9
 network:
-    arch: ViT-B/16  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True 
     drop_out: 0.0 
     emb_dropout: 0.0 

configs/hmdb51/hmdb_zero_shot.yaml

@@ -18,7 +18,7 @@ data:
         N: 0 #2
         M: 0  #9
 network:
-    arch: ViT-B/16  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True
     drop_out: 0.0
     emb_dropout: 0.0

configs/k400/k400_test.yaml

@@ -17,7 +17,7 @@ data:
     input_size: 224
     random_shift: False
 network:
-    arch: ViT-B/32  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True
     tsm: False
     drop_out: 0.0 

configs/k400/k400_train.yaml

@@ -20,7 +20,7 @@ data:
         M: 9  #9
     random_shift: True
 network:
-    arch: ViT-B/32  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True
     tsm: False
     drop_out: 0.0 

configs/k400/k400_zero_shot.yaml

@@ -16,7 +16,7 @@ data:
     input_size: 224
     random_shift: False
 network:
-    arch: ViT-B/32  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True
     tsm: False
     drop_out: 0.0 # probability of an element to be zeroed

configs/ucf101/ucf_test.yaml

@@ -18,7 +18,7 @@ data:
         N: 0 #2
         M: 0  #9
 network:
-    arch: ViT-B/16  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True 
     drop_out: 0.0 
     emb_dropout: 0.0 

configs/ucf101/ucf_train.yaml

@@ -21,7 +21,7 @@ data:
         N: 0 #2
         M: 0  #9
 network:
-    arch: ViT-B/16  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True  
     drop_out: 0.0 
     emb_dropout: 0.0 

configs/ucf101/ucf_zero_shot.yaml

@@ -18,7 +18,7 @@ data:
         N: 0 #2
         M: 0  #9
 network:
-    arch: ViT-B/16  #ViT-B/32 ViT-B/16
+    arch: CLIP-ViT-B/16  #CLIP-ViT-B/32 CLIP-ViT-B/16 TinyCLIP-ViT-61M/32-Text-29M TinyCLIP-ViT-40M/32-Text-19M
     init: True
     drop_out: 0.0
     emb_dropout: 0.0