diff --git a/clip/clip.py b/clip/clip.py index fbb9a4e..aefedb8 100644 --- a/clip/clip.py +++ b/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() diff --git a/clip/model.py b/clip/model.py index d92eda3..7b02461 100644 --- a/clip/model.py +++ b/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() \ No newline at end of file diff --git a/configs/hmdb51/hmdb_test.yaml b/configs/hmdb51/hmdb_test.yaml index 65d5c6a..09b840b 100644 --- a/configs/hmdb51/hmdb_test.yaml +++ b/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 diff --git a/configs/hmdb51/hmdb_train.yaml b/configs/hmdb51/hmdb_train.yaml index 5c31aa0..c19e710 100644 --- a/configs/hmdb51/hmdb_train.yaml +++ b/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 diff --git a/configs/hmdb51/hmdb_zero_shot.yaml b/configs/hmdb51/hmdb_zero_shot.yaml index 0be1b16..53693d9 100644 --- a/configs/hmdb51/hmdb_zero_shot.yaml +++ b/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 diff --git a/configs/k400/k400_test.yaml b/configs/k400/k400_test.yaml index bf92310..a9f448a 100644 --- a/configs/k400/k400_test.yaml +++ b/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 diff --git a/configs/k400/k400_train.yaml b/configs/k400/k400_train.yaml index d517345..5c6c639 100644 --- a/configs/k400/k400_train.yaml +++ b/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 diff --git a/configs/k400/k400_zero_shot.yaml b/configs/k400/k400_zero_shot.yaml index ab3d099..aa99183 100644 --- a/configs/k400/k400_zero_shot.yaml +++ b/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 diff --git a/configs/ucf101/ucf_test.yaml b/configs/ucf101/ucf_test.yaml index dcd05ea..c4b1f0c 100644 --- a/configs/ucf101/ucf_test.yaml +++ b/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 diff --git a/configs/ucf101/ucf_train.yaml b/configs/ucf101/ucf_train.yaml index 0626354..620c00d 100644 --- a/configs/ucf101/ucf_train.yaml +++ b/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 diff --git a/configs/ucf101/ucf_zero_shot.yaml b/configs/ucf101/ucf_zero_shot.yaml index d6b240e..ab05489 100644 --- a/configs/ucf101/ucf_zero_shot.yaml +++ b/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