|
import torch |
|
import torch.nn as nn |
|
from transformers import CLIPVisionModel |
|
|
|
|
|
class clip_vit_large_patch14_336(nn.Module): |
|
|
|
def __init__(self, vision_tower, use_resize_pos=True): |
|
super().__init__() |
|
|
|
self.is_loaded = False |
|
self.is_resize_pos = False |
|
|
|
self.vision_tower_name = vision_tower |
|
self.select_layer = -1 |
|
self.select_feature = 'patch' |
|
self.load_model() |
|
|
|
|
|
if use_resize_pos: |
|
self.resize_pos() |
|
|
|
def load_model(self): |
|
self.vision_tower = CLIPVisionModel.from_pretrained( |
|
self.vision_tower_name) |
|
self.vision_tower.requires_grad_(False) |
|
|
|
self.is_loaded = True |
|
|
|
def resize_pos(self): |
|
pos_embed_checkpoint = self.vision_tower.vision_model.embeddings.position_embedding.weight |
|
pos_embed_checkpoint = pos_embed_checkpoint.unsqueeze(0) |
|
orig_size = 24 |
|
new_size = 35 |
|
|
|
if pos_embed_checkpoint.shape[1] == new_size**2 + 1: |
|
self.is_resize_pos = True |
|
else: |
|
embedding_size = pos_embed_checkpoint.shape[-1] |
|
num_extra_tokens = 1 |
|
new_num = new_size**2 + num_extra_tokens |
|
|
|
|
|
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens] |
|
|
|
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:] |
|
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, |
|
embedding_size).permute( |
|
0, 3, 1, 2) |
|
pos_tokens = torch.nn.functional.interpolate( |
|
pos_tokens, |
|
size=(new_size, new_size), |
|
mode='bicubic', |
|
align_corners=False) |
|
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2) |
|
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1) |
|
|
|
new_pos_embed = new_pos_embed.squeeze(0) |
|
|
|
self.vision_tower.vision_model.embeddings.position_embedding = torch.nn.Embedding( |
|
new_num, 1024) |
|
self.vision_tower.vision_model.embeddings.position_embedding.weight = torch.nn.Parameter( |
|
new_pos_embed.to(pos_embed_checkpoint.dtype)) |
|
self.vision_tower.vision_model.embeddings.position_ids = torch.arange( |
|
new_num).expand((1, -1)) |
|
|
|
self.is_resize_pos = True |
|
|
|
def feature_select(self, image_forward_outs): |
|
image_features = image_forward_outs.hidden_states[self.select_layer] |
|
if self.select_feature == 'patch': |
|
image_features = image_features[:, 1:] |
|
elif self.select_feature == 'cls_patch': |
|
image_features = image_features |
|
else: |
|
raise ValueError( |
|
f'Unexpected select feature: {self.select_feature}') |
|
return image_features |
|
|
|
def forward(self, images): |
|
if not self.is_loaded: |
|
self.load_model() |
|
if type(images) is list: |
|
image_features = [] |
|
for image in images: |
|
image_forward_out = self.vision_tower( |
|
image.to(device=self.device, |
|
dtype=self.dtype).unsqueeze(0), |
|
output_hidden_states=True) |
|
image_feature = self.feature_select(image_forward_out).to( |
|
image.dtype) |
|
image_features.append(image_feature) |
|
else: |
|
image_forward_outs = self.vision_tower( |
|
images.to(device=self.device, dtype=self.dtype), |
|
output_hidden_states=True) |
|
image_features = self.feature_select(image_forward_outs).to( |
|
images.dtype) |
|
|
|
return image_features |
|
|
|
@property |
|
def device(self): |
|
return self.vision_tower.device |
|
|
|
@property |
|
def dtype(self): |
|
return self.vision_tower.dtype |
|
|
|
class DFN5B_CLIP_ViT_H_14_378(nn.Module): |
|
|
|
def __init__(self, vision_tower): |
|
super().__init__() |
|
|
|
self.is_loaded = False |
|
self.is_resize_pos = False |
|
|
|
self.vision_tower_name = vision_tower |
|
self.select_layer = -1 |
|
self.select_feature = 'patch' |
|
self.load_model() |
|
|
|
def load_model(self): |
|
self.vision_tower = CLIPVisionModel.from_pretrained( |
|
self.vision_tower_name) |
|
self.vision_tower.requires_grad_(False) |
|
|
|
self.is_loaded = True |
|
|
|
def feature_select(self, image_forward_outs): |
|
image_features = image_forward_outs.hidden_states[self.select_layer] |
|
if self.select_feature == 'patch': |
|
image_features = image_features[:, 1:] |
|
elif self.select_feature == 'cls_patch': |
|
image_features = image_features |
|
else: |
|
raise ValueError( |
|
f'Unexpected select feature: {self.select_feature}') |
|
return image_features |
|
|
|
def forward(self, images): |
|
if not self.is_loaded: |
|
self.load_model() |
|
if type(images) is list: |
|
image_features = [] |
|
for image in images: |
|
image_forward_out = self.vision_tower( |
|
image.to(device=self.device, |
|
dtype=self.dtype).unsqueeze(0), |
|
output_hidden_states=True) |
|
image_feature = self.feature_select(image_forward_out).to( |
|
image.dtype) |
|
image_features.append(image_feature) |
|
else: |
|
image_forward_outs = self.vision_tower( |
|
images.to(device=self.device, dtype=self.dtype), |
|
output_hidden_states=True) |
|
image_features = self.feature_select(image_forward_outs).to( |
|
images.dtype) |
|
|
|
return image_features |
|
|
|
@property |
|
def device(self): |
|
return self.vision_tower.device |
|
|
|
@property |
|
def dtype(self): |
|
return self.vision_tower.dtype |