Delete vision_tower.py

vision_tower.py

@@ -1,161 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.init as init
-import torch.nn.functional as F
-
-from paths import *
-
-from typing import Dict, List, Optional, Set, Tuple, Union
-from transformers import AutoImageProcessor, AutoModel, Dinov2Model
-from transformers.models.dinov2.modeling_dinov2 import Dinov2Embeddings
-from transformers.models.dinov2.configuration_dinov2 import Dinov2Config
-import numpy as np
-from contextlib import nullcontext
-
-def get_activation(activation):
-    if activation.lower() == 'gelu':
-        return nn.GELU()
-    elif activation.lower() == 'rrelu':
-        return nn.RReLU(inplace=True)
-    elif activation.lower() == 'selu':
-        return nn.SELU(inplace=True)
-    elif activation.lower() == 'silu':
-        return nn.SiLU(inplace=True)
-    elif activation.lower() == 'hardswish':
-        return nn.Hardswish(inplace=True)
-    elif activation.lower() == 'leakyrelu':
-        return nn.LeakyReLU(inplace=True)
-    elif activation.lower() == 'sigmoid':
-        return nn.Sigmoid()
-    elif activation.lower() == 'tanh':
-        return nn.Tanh()
-    else:
-        return nn.ReLU(inplace=True)
-
-
-
-class MLP_dim(nn.Module):
-    def __init__(
-        self, in_dim=512, out_dim=1024, bias=True, activation='relu'):
-        super().__init__()
-        self.act = get_activation(activation)
-        self.net1 = nn.Sequential(
-            nn.Linear(in_dim, int(out_dim), bias=bias),
-            nn.BatchNorm1d(int(out_dim)),
-            self.act
-        )
-        self.net2 = nn.Sequential(
-            nn.Linear(int(out_dim), out_dim, bias=bias),
-            nn.BatchNorm1d(out_dim)
-        )
-
-    def forward(self, x):
-        return self.net2(self.net1(x))
-
-class FLIP_Dinov2Embeddings(Dinov2Embeddings):
-    """
-    Construct the CLS token, mask token, position and patch embeddings.
-    """
-
-    def __init__(self, config: Dinov2Config) -> None:
-        super().__init__(config)
-
-    def forward(self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.Tensor] = None) -> torch.Tensor:
-        batch_size, _, height, width = pixel_values.shape
-        target_dtype = self.patch_embeddings.projection.weight.dtype
-        embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))
-
-        # add the [CLS] token to the embedded patch tokens
-        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
-        embeddings = torch.cat((cls_tokens, embeddings), dim=1)
-
-        # add positional encoding to each token
-        embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
-
-        if bool_masked_pos is not None:
-            # embeddings = torch.where(
-            #     bool_masked_pos.unsqueeze(-1), self.mask_token.to(embeddings.dtype).unsqueeze(0), embeddings
-            # )
-            B,S,D = embeddings.shape
-            batch_indices = torch.arange(B).unsqueeze(1)
-            embeddings = embeddings[batch_indices, bool_masked_pos]
-
-        embeddings = self.dropout(embeddings)
-
-        return embeddings
-
-class FLIP_DINOv2(Dinov2Model):
-    def __init__(self, config):
-        super().__init__(config)
-        
-        self.embeddings = FLIP_Dinov2Embeddings(config)
-        
-class DINOv2_MLP(nn.Module):
-    def __init__(self, 
-                 dino_mode,
-                 in_dim,
-                 out_dim,
-                 evaluate,
-                 mask_dino,
-                 frozen_back
-                ) -> None:
-        super().__init__()
-        # self.dinov2 = AutoModel.from_pretrained(DINO_BASE)
-        if dino_mode == 'base':
-            self.dinov2 = FLIP_DINOv2.from_pretrained(DINO_BASE, cache_dir='./')
-        elif dino_mode == 'large':
-            self.dinov2 = FLIP_DINOv2.from_pretrained(DINO_LARGE, cache_dir='./')
-        elif dino_mode == 'small':
-            self.dinov2 = FLIP_DINOv2.from_pretrained(DINO_SMALL, cache_dir='./')
-        elif dino_mode == 'giant':
-            self.dinov2 = FLIP_DINOv2.from_pretrained(DINO_GIANT, cache_dir='./')
-        
-        self.down_sampler = MLP_dim(in_dim=in_dim, out_dim=out_dim)
-        self.random_mask  = False
-        if not evaluate:
-            self.init_weights(self.down_sampler)
-            self.random_mask = mask_dino
-        if frozen_back:
-            self.forward_mode = torch.no_grad()
-        else:
-            self.forward_mode = nullcontext()
-        
-    def forward(self, img_inputs):
-        device = self.get_device()
-        # print(img_inputs['pixel_values'].shape)
-        
-        with self.forward_mode:
-            if self.random_mask:
-                B = len(img_inputs['pixel_values'])
-                S = 256
-                indices = []
-                for i in range(B):
-                    tmp = torch.randperm(S)[:S//2]
-                    tmp = tmp.sort().values + 1
-                    indices.append(tmp)
-                indices = torch.stack(indices, dim=0)
-                indices = torch.cat([torch.zeros(B, 1, dtype=torch.long, device='cpu'), indices], dim=1)
-                # print(indices.shape)
-                img_inputs['bool_masked_pos'] = indices.to(device)
-                
-            dino_outputs = self.dinov2(**img_inputs)
-            dino_seq = dino_outputs.last_hidden_state
-            # B,S,_ = dino_seq.shape
-            # dino_seq = dino_seq.view(B*S,-1)
-            dino_seq = dino_seq[:,0,:]
-        
-        down_sample_out = self.down_sampler(dino_seq)
-        # down_sample_out = down_sample_out.view(B,S,-1)
-        # down_sample_out = down_sample_out[:,0,:]
-        
-        return down_sample_out
-    
-    def get_device(self):
-        return next(self.parameters()).device
-    
-    def init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            init.xavier_uniform_(m.weight)
-            if m.bias is not None:
-                init.constant_(m.bias, 0)
-