upload

alignedthreeattn_model.py

@@ -20,6 +20,28 @@ class ThreeAttnNodes(nn.Module):
         for backbone in [self.backbone1, self.backbone2, self.backbone3]:
             backbone.requires_grad_(False)
             backbone.eval()
+                    
+        def resample_position_embeddings(embeddings, h, w):
+            cls_embeddings = embeddings[0]
+            patch_embeddings = embeddings[1:]  # [14*14, 768]
+            hw = np.sqrt(patch_embeddings.shape[0]).astype(int)
+            patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=hw)
+            patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(h, w), mode="nearest").squeeze(0)
+            patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
+            embeddings = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
+            return embeddings
+
+        pos_embd = self.backbone1.model.visual.positional_embedding 
+        pos_embd = resample_position_embeddings(pos_embd, 42, 42)
+        self.backbone1.model.visual.positional_embedding = nn.Parameter(pos_embd)
+
+        pos_embed = self.backbone3.model.pos_embed[0]
+        pos_embed = resample_position_embeddings(pos_embed, 42, 42)
+        self.backbone3.model.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
+        self.backbone3.model.img_size = (672, 672)
+        self.backbone3.model.patch_embed.img_size = (672, 672)
+            
+            
         self.align_weights = align_weights
             
     @torch.no_grad()