Update MT.py

MT.py

@@ -131,7 +131,7 @@ class PositionEmbeddingSine(nn.Module):
         return pos
 
 
-def feature_add_position(feature0, feature_channels, scale=1.0):
+def feature_add_position(feature0, feature_channels, scale=0.5):
     temp = torch.mean(abs(feature0))
     pos_enc = PositionEmbeddingSine(num_pos_feats=feature_channels // 2)
     # position = PositionalEncodingPermute2D(feature_channels)(feature0)
@@ -223,8 +223,6 @@ class TransformerLayer(nn.Module):
             att = feature_add_position(att.transpose(-1, -2).view(
                 B, C, shape[0], shape[1]), C).reshape(B, C, -1).transpose(-1, -2)
 
-            # att = feature_add_position(att.transpose(-1, -2).view(
-            #     B, C, shape[0], shape[1]), C).reshape(B, C, -1).transpose(-1, -2)
         val_proj = self.v_proj(value)
         att_proj = self.att_proj(att)  # [B, L, C]
         norm_fac = torch.sum(att_proj ** 2, dim=-1, keepdim=True) ** 0.5
@@ -237,7 +235,6 @@ class TransformerLayer(nn.Module):
         D = 1 / (torch.sqrt(D) + 1e-6)  # normalized node degrees
         A = D * A * D.transpose(-1, -2)
 
-        # A = torch.softmax(A , dim=2)  # [B, L, L]
         message = torch.matmul(A, val_proj)  # [B, L, C]
 
         message = self.merge(message)  # [B, L, C]
@@ -246,9 +243,6 @@ class TransformerLayer(nn.Module):
             message = self.mlp(torch.cat([value, message], dim=-1))
             message = self.norm2(message)
 
-        # if iteration > 2:
-        #     message = self.drop(message)
-
         att = self.attn_updater(att, message, shape)
         value = self.gru(value, message, shape)
         return value, att, A
@@ -290,14 +284,11 @@ class FeatureTransformer(nn.Module):
         att = att.flatten(-2).permute(0, 2, 1)  # [B, H*W, C]
         for i in range(self.num_layers):
             value, att, attn_viz = self.layers(att=att, value=value, shape=[h, w], iteration=i)
-            attn_viz = attn_viz.reshape(b, h, w, h, w)
-            attn_viz_list.append(attn_viz)
-            value_decode = self.normalize(
-                torch.square(self.re_proj(value)))  # map to motion energy, Do use normalization here
-            # print("value_decode",value_decode.abs().mean())
-            attn_list.append(att.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
-            feature_list.append(value_decode.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
-        # reshape back
+            value_decode = self.normalize(torch.square(self.re_proj(value)))  # map to motion energy, Do use normalization here
+        
+        attn_viz_list.append(attn_viz.reshape(b, h, w, h, w))
+        attn_list.append(att.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
+        feature_list.append(value_decode.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
         return feature_list, attn_list, attn_viz_list
 
     def forward_save_mem(self, feature0, add_position_embedding=True):