Fixed error

src/dinotext.py → dinotext.py

@@ -16,14 +16,14 @@ from transformers import BertModel, AutoTokenizer
 import torchvision.transforms as T
 import clip
 import importlib
-import src.us as us
+from .us import normalize
 
-from src.pamr import PAMR
-from src.masker import DINOTextMasker
-from src.templates import get_template
+from .pamr import PAMR
+from .masker import DINOTextMasker
+from .templates import get_template
 
-from src.model import ProjectionLayer, VisualProjectionLayer, CLIPLastLayer, DoubleMLP
-from src.hooks import average_text_tokens, get_vit_out, feats
+from .model import ProjectionLayer, VisualProjectionLayer, CLIPLastLayer, DoubleMLP
+from .hooks import average_text_tokens, get_vit_out, feats
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -280,7 +280,7 @@ class DINOText(nn.Module):
         text_embs = text_embs.mean(dim=1).float()
         if type(self.proj) == ProjectionLayer or type(self.proj) == DoubleMLP:
             text_embs = self.proj.project_clip_txt(text_embs)
-        text_embs = us.normalize(text_embs, dim=-1)
+        text_embs = normalize(text_embs, dim=-1)
 
         return text_embs
 

src/masker.py → masker.py

@@ -8,11 +8,11 @@ import torch
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
-import src.us as us
+from .us import normalize
 from einops import rearrange, repeat
 
 # from models.dinotext.gumbel import gumbel_sigmoid
-from src.modules import FeatureEncoder
+from .modules import FeatureEncoder
 from omegaconf import OmegaConf
 
 
@@ -129,14 +129,14 @@ class Masker(nn.Module):
         B = image.size(0)
         image_emb, feats = self.image_encoder(image, image_feat, ret_feats=True)  # [BCHW]
 
-        image_emb_norm = us.normalize(image_emb, dim=1)
-        text_emb_norm = us.normalize(text_emb, dim=-1)
+        image_emb_norm = normalize(image_emb, dim=1)
+        text_emb_norm = normalize(text_emb, dim=-1)
 
         H, W = image_emb.shape[2:]
         D = dist.get_world_size()
 
         # simmap [B, B*D, H, W] where D is #devices
-        all_text_emb_norm = us.gather_cat(text_emb_norm, grad=True, contiguous_grad=True)
+        all_text_emb_norm = gather_cat(text_emb_norm, grad=True, contiguous_grad=True)
         simmap = torch.einsum("bchw,nc->bnhw", image_emb_norm, all_text_emb_norm)
         mask, soft_mask = self.sim2mask(simmap, deterministic=deterministic)
 
@@ -178,8 +178,8 @@ class Masker(nn.Module):
         """
         image_emb = self.image_encoder(image, image_feat)  # [BCHW]
 
-        image_emb = us.normalize(image_emb, dim=1)  # BCHW
-        text_emb = us.normalize(text_emb, dim=-1)  # NC
+        image_emb = normalize(image_emb, dim=1)  # BCHW
+        text_emb = normalize(text_emb, dim=-1)  # NC
 
         simmap = torch.einsum("b c h w, n c -> b n h w", image_emb, text_emb)
 
@@ -219,8 +219,8 @@ class DINOTextMasker(nn.Module):
         n, c = text_emb.shape
 
         if self.similarity_type == "cosine":
-            image_feat = us.normalize(image_feat, dim=1)  # BCHW
-            # text_emb = us.normalize(text_emb, dim=-1)  # NKC
+            image_feat = normalize(image_feat, dim=1)  # BCHW
+            # text_emb = normalize(text_emb, dim=-1)  # NKC
             simmap = torch.einsum("b c h w, n c -> b n h w", image_feat, text_emb)
         else:
             raise NotImplementedError("similarity type {} not implemented".format(self.similarity_type))

src/model.py → model.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from src.hooks import get_self_attention, process_self_attention, feats
+from .hooks import get_self_attention, process_self_attention, feats
 
 class VisualProjectionLayer(nn.Module):
     """

modeling_talk2dino.py

@@ -1,6 +1,6 @@
-from src.dinotext import DINOText
+from .configuration_talk2dino import Talk2DINOConfig
+from .dinotext import DINOText
 from transformers import PreTrainedModel
-from configuration_talk2dino import Talk2DINOConfig
 import clip
 import torch