File size: 3,214 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import torch
import random
import logging
import open_clip
import torch.nn as nn
import src.clip.clip as CLIP
def load_model(architecture, **kargs):
# architecture parameter split
# e.g: ViT-B/16|laion2b_s34b_b79k
params = architecture.split('|')
clip_arch = params[0]
if (len(params) == 1):
model, transform = CLIP.load(
clip_arch,
"cpu",
**kargs
)
elif (len(params) == 2):
open_pretrain = params[1]
_model, _, _ = open_clip.create_model_and_transforms(
clip_arch.replace("/", "-"),
pretrained=open_pretrain,
device="cpu"
)
model, transform = CLIP.load(
_model.state_dict(),
"cpu",
**kargs
)
del _model
elif (len(params) > 2):
raise NotImplementedError()
return model, transform
class VideoAttrExtractor(nn.Module):
def __init__(
self,
architecture,
text_embed,
store_attrs=[],
attn_record=False,
pretrain=None
):
super().__init__()
self.model, self.transform = load_model(
architecture,
store_attrs=store_attrs,
attn_record=attn_record
)
self.model = self.model.visual.float()
if (pretrain):
logging.info("Loading image encoder pretrain weights...")
state_dict = torch.load(pretrain, "cpu")
try:
self.model.load_state_dict(state_dict, strict=True)
except:
conflicts = self.model.load_state_dict(state_dict, strict=False)
logging.warning(
f"during visual pretrain weights loading, disabling strict mode with conflicts:\n{conflicts}"
)
self.model.requires_grad_(False)
if not text_embed:
self.model.proj = None
self.feat_dim = self.model.transformer.width
else:
self.feat_dim = self.model.output_dim
@property
def n_px(self):
return self.model.input_resolution
@property
def n_layers(self):
return self.model.transformer.layers
@property
def n_heads(self):
return self.model.transformer.heads
@property
def patch_size(self):
return self.model.patch_size
@property
def patch_num(self):
return self.model.patch_num
@property
def n_patch(self):
return int(self.n_px // self.patch_size)
@property
def embed_dim(self):
return self.feat_dim
def forward(self, x):
b, t = x.shape[:2]
# pass throught for attributes
embeds = self.model(x)
# retrieve all layer attributes
layer_attrs = []
for blk in self.model.transformer.resblocks:
attrs = blk.pop_attr()
layer_attrs.append(attrs)
return dict(
layer_attrs=layer_attrs,
embeds=embeds
)
def train(self, mode=True):
self.model.eval()
return self
if __name__ == "__main__":
VideoAttrExtractor(
"ViT-L/14|laion2b_s32b_b82k",
text_embed=False
)
|