Upload DINOSAUR.py

slots/DINOSAUR.py

@@ -221,7 +221,78 @@ class Decoder(nn.Module):
 
         slot_maps = self.layer4(slot_maps)  #  (B * S, token, 1024 + 1)
 
-        return slot_maps
+        return slot_maps, slot_maps
+
+
+class Decoder_to_DINOV2(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        # === Token calculations ===
+        slot_dim = args['slot_dim']
+        hidden_dim = 2048
+
+        # === MLP Based Decoder ===
+        self.layer1 = nn.Linear(slot_dim, hidden_dim)
+        self.layer2 = nn.Linear(hidden_dim, hidden_dim)
+        self.layer3 = nn.Linear(hidden_dim, hidden_dim)
+        self.layer4 = nn.Linear(hidden_dim, 1024 + 1)
+
+        self.layer_to_dinov2 = nn.Linear(hidden_dim, 768)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, slot_maps):
+        # :arg slot_maps: (B * S, token, D_slot)
+        slot_maps =  self.relu(self.layer1(slot_maps))  #  (B * S, token, D_hidden)
+        x_dinov2 = self.layer_to_dinov2(slot_maps)
+        slot_maps = self.relu(self.layer2(slot_maps))  #  (B * S, token, D_hidden)
+        slot_maps = self.relu(self.layer3(slot_maps))  #  (B * S, token, D_hidden)
+
+        slot_maps = self.layer4(slot_maps)  #  (B * S, token, 1024 + 1)
+
+        return slot_maps, x_dinov2
+
+from torch.nn.init import trunc_normal_
+class DINOHead(nn.Module):
+    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, nlayers=3, hidden_dim=2048, bottleneck_dim=768):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        if nlayers == 1:
+            self.mlp = nn.Linear(in_dim, bottleneck_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+            for _ in range(nlayers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if use_bn:
+                    layers.append(nn.BatchNorm1d(hidden_dim))
+                layers.append(nn.GELU())
+            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            self.mlp = nn.Sequential(*layers)
+        self.apply(self._init_weights)
+        self.gelu = nn.GELU()
+        self.last_layer1 = nn.Linear(bottleneck_dim, bottleneck_dim)
+        self.last_layer2 = nn.Linear(bottleneck_dim, out_dim)
+
+        # self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+        # self.last_layer.weight_g.data.fill_(1)
+        # if norm_last_layer:
+        #     self.last_layer.weight_g.requires_grad = False
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x_dinov2 = self.mlp(x)
+        # x = nn.functional.normalize(x, dim=-1, p=2)
+        x = self.gelu(self.last_layer1(x_dinov2))
+        x = self.last_layer2(x)
+        return x, x_dinov2
 
 
 class ISA(nn.Module):
@@ -517,7 +588,9 @@ class DINOSAURpp(nn.Module):
         else:
             self.slot_encoder = SA(args, input_dim=1024)
 
-        self.slot_decoder = Decoder(args)
+        self.slot_decoder = Decoder(args) #ori easy mlp
+        # self.slot_decoder = DINOHead(in_dim=256, out_dim=1024+1, nlayers=3, bottleneck_dim=768) #ori easy mlp
+        # self.slot_decoder = Decoder_to_DINOV2(args) #ori easy mlp
 
         self.pos_dec = nn.Parameter(torch.Tensor(1, self.token_num, self.slot_dim))
         init.normal_(self.pos_dec, mean=0., std=.02)
@@ -572,17 +645,17 @@ class DINOSAURpp(nn.Module):
             rel_grid = self.slot_encoder.get_rel_grid(attn)  # (B, S, token, D_slot)
 
             slot_maps = self.sbd_slots(slots) + rel_grid  # (B, S, token, D_slot)
-            slot_maps = self.slot_decoder(slot_maps)  # (B, S, token, 1024 + 1)
+            slot_maps, x_dinov2 = self.slot_decoder(slot_maps)  # (B, S, token, 1024 + 1)
 
         else:
             slots = self.slot_encoder(features)  # (B, S, D_slot), (B, S, token)
             assert torch.sum(torch.isnan(slots)) == 0
 
             slot_maps, pos_maps = self.sbd_slots(slots)
-            slot_maps = self.slot_decoder(slot_maps)  # (B, S, token, 1024 + 1)
+            slot_maps, x_dinov2 = self.slot_decoder(slot_maps)  # (B, S, token, 1024 + 1)
 
         reconstruction, masks = self.reconstruct_feature_map(slot_maps)  # (B, token, 1024), (B, S, token)
 
-        return reconstruction, slots, masks
+        return reconstruction, slots, masks, x_dinov2