Upload model

modeling_solider.py

@@ -1234,7 +1234,9 @@ class SwinTransformer(BaseModule):
         convert_weights=False,
         frozen_stages=-1,
         init_cfg=None,
-        semantic_weight=0.0,
+        # NOTE: This is my modification based on SOLIDER
+        semantic_weight=0.5,
+        freeze_semantic_embedding=False,
     ):
         self.convert_weights = convert_weights
         self.frozen_stages = frozen_stages
@@ -1341,6 +1343,7 @@ class SwinTransformer(BaseModule):
 
         # semantic embedding
         self.semantic_weight = semantic_weight
+        self.freeze_semantic_embedding = freeze_semantic_embedding
         if self.semantic_weight >= 0:
             self.semantic_embed_w = ModuleList()
             self.semantic_embed_b = ModuleList()
@@ -1350,10 +1353,11 @@ class SwinTransformer(BaseModule):
                 semantic_embed_w = nn.Linear(2, self.num_features[i + 1])
                 semantic_embed_b = nn.Linear(2, self.num_features[i + 1])
                 # TODO: Test with semantic embed unfreeze
-                for param in semantic_embed_w.parameters():
-                    param.requires_grad = False
-                for param in semantic_embed_b.parameters():
-                    param.requires_grad = False
+                if self.freeze_semantic_embedding:
+                    for param in semantic_embed_w.parameters():
+                        param.requires_grad = False
+                    for param in semantic_embed_b.parameters():
+                        param.requires_grad = False
                 trunc_normal_init(semantic_embed_w, std=0.02, bias=0.0)
                 trunc_normal_init(semantic_embed_b, std=0.02, bias=0.0)
                 self.semantic_embed_w.append(semantic_embed_w)