update

build_model.py

@@ -0,0 +1,80 @@
+import argparse
+from typing import Any, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_
+from copy import deepcopy
+import os
+import torch.backends.cudnn as cudnn
+
+import models.vision_transformer as vits
+
+class vit(nn.Module):
+    
+    def __init__(self, model_size="base", freeze_transformer=True, pretrained_weights=None):
+        super(ibotvit, self).__init__()
+        self.model_size = model_size
+        self.freeze_transformer = freeze_transformer
+        self.pretrained_weights = pretrained_weights
+
+        # Loading a model with registers
+        n_register_tokens = 4
+        
+        if model_size == "vit_small":
+            self.embedding_size = 384
+            
+        elif model_size == "vit_base":
+            self.embedding_size = 768
+
+        elif model_size == "vit_large":
+            self.embedding_size = 1024
+            
+        elif model_size == "giant":
+            self.embedding_size = 1536
+
+        # Load state_dict
+        model = vits.__dict__[model_size](patch_size=16)
+        self.transformer = deepcopy(model)
+
+        # Freeze transformer if specified        
+        if self.freeze_transformer:
+            for param in self.transformer.parameters():
+                param.requires_grad = False
+
+        
+        if self.pretrained_weights and os.path.isfile(self.pretrained_weights):
+            state_dict = torch.load(self.pretrained_weights, map_location="cpu")
+            if 'teacher' in state_dict:
+                state_dict = state_dict['teacher']
+            elif 'model' in state_dict:
+                state_dict = state_dict['model']
+
+            # remove `backbone.` prefix induced by multicrop wrapper
+            state_dict = {
+                (k[len("teacher."):] if k.startswith("teacher.") else k): v
+                for k, v in state_dict.items()
+            }
+            state_dict = {
+                (k[len("backbone."):] if k.startswith("backbone.") else k): v
+                for k, v in state_dict.items()
+            }
+            msg = self.transformer.load_state_dict(state_dict, strict=False)
+            print(model_size, msg)
+        
+
+    def forward(self, x):
+        x = self.transformer(x)
+
+        return x
+
+
+
+def build_model(args):
+    
+    net = vit("vit_base", freeze_transformer=True, pretrained_weights=args.pretrained_weights)
+    net.cuda()
+    
+
+
+    return net

models/build_model.py

@@ -0,0 +1,80 @@
+import argparse
+from typing import Any, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_
+from copy import deepcopy
+import os
+import torch.backends.cudnn as cudnn
+
+import models.vision_transformer as vits
+
+class vit(nn.Module):
+    
+    def __init__(self, model_size="base", freeze_transformer=True, pretrained_weights=None):
+        super(ibotvit, self).__init__()
+        self.model_size = model_size
+        self.freeze_transformer = freeze_transformer
+        self.pretrained_weights = pretrained_weights
+
+        # Loading a model with registers
+        n_register_tokens = 4
+        
+        if model_size == "vit_small":
+            self.embedding_size = 384
+            
+        elif model_size == "vit_base":
+            self.embedding_size = 768
+
+        elif model_size == "vit_large":
+            self.embedding_size = 1024
+            
+        elif model_size == "giant":
+            self.embedding_size = 1536
+
+        # Load state_dict
+        model = vits.__dict__[model_size](patch_size=16)
+        self.transformer = deepcopy(model)
+
+        # Freeze transformer if specified        
+        if self.freeze_transformer:
+            for param in self.transformer.parameters():
+                param.requires_grad = False
+
+        
+        if self.pretrained_weights and os.path.isfile(self.pretrained_weights):
+            state_dict = torch.load(self.pretrained_weights, map_location="cpu")
+            if 'teacher' in state_dict:
+                state_dict = state_dict['teacher']
+            elif 'model' in state_dict:
+                state_dict = state_dict['model']
+
+            # remove `backbone.` prefix induced by multicrop wrapper
+            state_dict = {
+                (k[len("teacher."):] if k.startswith("teacher.") else k): v
+                for k, v in state_dict.items()
+            }
+            state_dict = {
+                (k[len("backbone."):] if k.startswith("backbone.") else k): v
+                for k, v in state_dict.items()
+            }
+            msg = self.transformer.load_state_dict(state_dict, strict=False)
+            print(model_size, msg)
+        
+
+    def forward(self, x):
+        x = self.transformer(x)
+
+        return x
+
+
+
+def build_model(args):
+    
+    net = vit("vit_base", freeze_transformer=True, pretrained_weights=args.pretrained_weights)
+    net.cuda()
+    
+
+
+    return net