Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torchmetrics.functional import retrieval_average_precision | |
| import pytorch_lightning as pl | |
| from src.dinov2.models.vision_transformer import vit_base | |
| from functools import partial | |
| # from src.clip import clip | |
| from src.options import opts | |
| def freeze_model(m): | |
| m.requires_grad_(False) | |
| def freeze_all_but_bn(m): | |
| if not isinstance(m, torch.nn.LayerNorm): | |
| if hasattr(m, 'weight') and m.weight is not None: | |
| m.weight.requires_grad_(False) | |
| if hasattr(m, 'bias') and m.bias is not None: | |
| m.bias.requires_grad_(False) | |
| else: | |
| print("LayerNorm") | |
| class Model(pl.LightningModule): | |
| def __init__(self): | |
| super().__init__() | |
| self.opts = opts | |
| self.dino = vit_base(patch_size=14, block_chunks=0, init_values=1.0) | |
| print("self.dino", self.dino) | |
| # Prompt Engineering | |
| self.sk_prompt = nn.Parameter(torch.randn(self.opts.n_prompts, self.opts.prompt_dim)) | |
| self.img_prompt = nn.Parameter(torch.randn(self.opts.n_prompts, self.opts.prompt_dim)) | |
| self.distance_fn = lambda x, y: 1.0 - F.cosine_similarity(x, y) | |
| self.loss_fn_triplet = nn.TripletMarginWithDistanceLoss( | |
| distance_function=self.distance_fn, margin=0.2) | |
| self.emb_cos_loss = nn.CosineEmbeddingLoss(margin=0.2) | |
| self.loss_kl = nn.KLDivLoss(reduction="batchmean", log_target=True) | |
| self.best_metric = -1e3 | |
| # normalization layer for the representations z1 and z2 | |
| # self.bn = nn.BatchNorm1d(self.opts.prompt_dim, affine=False) | |
| def configure_optimizers(self): | |
| if self.opts.model_type == 'one_encoder': | |
| model_params = list(self.dino.parameters()) | |
| else: | |
| model_params = list(self.dino.parameters()) + list(self.clip_sk.parameters()) | |
| optimizer = torch.optim.Adam([ | |
| {'params': model_params, 'lr': self.opts.clip_LN_lr}, | |
| {'params': [self.sk_prompt] + [self.img_prompt], 'lr': self.opts.prompt_lr}]) | |
| return optimizer | |
| def forward(self, data, dtype='image'): | |
| if dtype == 'image': | |
| feat = self.dino(data, prompt=self.img_prompt.expand(data.shape[0], -1, -1)) | |
| else: | |
| feat = self.dino(data, prompt=self.sk_prompt.expand(data.shape[0], -1, -1)) | |
| return feat |