Update mmseg/ttp/models.py

mmseg/ttp/models.py

@@ -52,6 +52,13 @@ class MMPretrainSamVisionEncoder(BaseModule):
             # freeze the vision encoder
             for param in self.vision_encoder.parameters():
                 param.requires_grad = False
+        for name, param in self.vision_encoder.named_parameters():
+            if 'down_channel' in name:
+                param.requires_grad = True
+            if 'soft_ffn' in name:
+                param.requires_grad = True
+        if is_main_process() and peft_cfg is not None:
+            self.vision_encoder.print_trainable_parameters()
 
     def forward(self, x):
         return self.vision_encoder(x)
@@ -265,7 +272,7 @@ class TimeFusionTransformerEncoderLayer(BaseModule):
             dropout_layer=dict(type='DropPath', drop_prob=drop_path_rate),
             act_cfg=act_cfg)
 
-        if self.window_size > 0:
+        if self.window_size == 0:
             in_channels = embed_dims * 2
             self.down_channel = nn.Conv2d(in_channels, 1, kernel_size=1, stride=1, bias=False)
             self.down_channel.weight.data.fill_(1.0/in_channels)
@@ -300,7 +307,7 @@ class TimeFusionTransformerEncoderLayer(BaseModule):
 
         x = self.ffn(self.ln2(x), identity=x)
         # # time phase fusion
-        if self.window_size > 0:
+        if self.window_size == 0:
             x = einops.rearrange(x, 'b h w d -> b d h w')  # 2B, C, H, W
             x0 = x[:x.size(0)//2]
             x1 = x[x.size(0)//2:]  # B, C, H, W