MATH-LLM-7B / Vision_Tower.py
ALmonster's picture
Upload 17 files
9487267 verified
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()
#change model to input shape[490*490]
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 #336/14
new_size = 35 #490/14
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
#print('Position interpolate from %dx%d to %dx%d' %
# (orig_size, orig_size, new_size, new_size))
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
# only the position tokens are interpolated
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: # not batch infurence speed!
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: # not batch infurence speed!
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