| import torch | |
| import torch.nn.functional as F | |
| import torch.nn as nn | |
| from open_clip.model import _build_vision_tower | |
| class CLIP(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| model_name = 'convnext_large' | |
| vision_cfg = {'timm_model_name': model_name, 'timm_model_pretrained': False, 'timm_pool': '', 'timm_proj': 'mlp', 'timm_drop': 0.0, 'timm_drop_path': 0.1, 'image_size': 320} | |
| self.visual = _build_vision_tower(embed_dim=768, vision_cfg=vision_cfg, quick_gelu=False) | |
| self.eval() | |
| self.freeze_everything() | |
| def freeze_everything(self): | |
| for param in self.visual.parameters(): | |
| param.requires_grad = False | |
| def extract_features(self, x): | |
| out = {} | |
| x = x.to(self.visual.trunk.stem.state_dict()['1.bias'].dtype) | |
| x = self.visual.trunk.stem(x) | |
| out['stem'] = x.contiguous() | |
| for i in range(4): | |
| x = self.visual.trunk.stages[i](x) | |
| out[f'res{i+2}'] = x.contiguous() | |
| x = self.visual.trunk.norm_pre(x) | |
| out['clip_vis_dense'] = x.contiguous() | |
| return out | |
| def forward(self, x): | |
| self.eval() | |
| with torch.no_grad(): | |
| return self.extract_features(x) | |