Fix SSF implementation

micro_sam/models/peft_sam.py

@@ -112,10 +112,10 @@ def __init__(self, layer, dim):
         self.layer = layer
         self.scale = nn.Parameter(torch.normal(mean=1.0, std=0.2, size=(dim,)))
         self.shift = nn.Parameter(torch.normal(mean=0.0, std=0.2, size=(dim,)))
+        layer = self
 
     def forward(self, x):
         x = self.layer(x)
-
         assert self.scale.shape == self.shift.shape
         if x.shape[-1] == self.scale.shape[0]:
             return x * self.scale + self.shift
@@ -133,37 +133,25 @@ class SSFSurgery(nn.Module):
         block: The chosen attention blocks for implementing ssf.
         dim: The input dimensions determining the shape of scale and shift parameters.
     """
-    def __init__(self, rank: int, block: nn.Module, dim: Optional[int] = None):
+    def __init__(self, rank: int, block: nn.Module):
         super().__init__()
         self.block = block
 
         # If we get a transformer block (w. multiple sub-layers), we perform surgery on each layer.
         if hasattr(block, "attn"):  # the minimum assumption is to verify the attention layers.
-            self.scale_shift_layers = nn.ModuleList(self.add_scale_shift_layers_to_block(block))
-        else:  # This is an individual layer after which we apply scale and shift.
-            if dim is None:
-                raise ValueError("'dim' must be provided for the scale and shift parameters.")
-            self.scale_shift_layers = nn.ModuleList([self.create_scale_shift_layer(layer=block, dim=dim)])
-
-    def add_scale_shift_layers_to_block(self, block):
-        peft_blocks = [
-            ScaleShiftLayer(block.attn.qkv, block.attn.qkv.in_features),
-            ScaleShiftLayer(block.attn.proj, block.attn.proj.in_features),
-            ScaleShiftLayer(block.mlp.lin1, block.mlp.lin1.in_features),
-            ScaleShiftLayer(block.mlp.lin2, block.mlp.lin2.in_features),
-            ScaleShiftLayer(block.norm1, block.norm1.normalized_shape[0]),
-            ScaleShiftLayer(block.norm2, block.norm2.normalized_shape[0]),
-        ]
-        return nn.ModuleList(peft_blocks)
-
-    def create_scale_shift_layer(self, layer, dim):
-        return ScaleShiftLayer(layer=layer, dim=dim)
+            block.attn.qkv = ScaleShiftLayer(block.attn.qkv, block.attn.qkv.in_features*3)
+            block.attn.proj = ScaleShiftLayer(block.attn.proj, block.attn.proj.in_features)
+            block.mlp.lin1 = ScaleShiftLayer(block.mlp.lin1, block.mlp.lin1.out_features)
+            block.mlp.lin2 = ScaleShiftLayer(block.mlp.lin2, block.mlp.lin2.out_features)
+            block.norm1 = ScaleShiftLayer(block.norm1, block.norm1.normalized_shape[0])
+            block.norm2 = ScaleShiftLayer(block.norm2, block.norm2.normalized_shape[0])
 
-    def forward(self, x):
-        for layer in self.scale_shift_layers:
-            x = layer(x)
+        # If we get the embedding block, add one ScaleShiftLayer
+        elif hasattr(block, "patch_embed"):
+            block.proj = ScaleShiftLayer(block.proj, block.proj.out_channels)
 
-        return self.block(x)
+    def forward(self, x):
+        return x
 
 
 class SelectiveSurgery(nn.Module):
@@ -269,8 +257,7 @@ def __init__(
         if issubclass(self.peft_module, SSFSurgery):
             self.peft_blocks.append(
                 self.peft_module(
-                    rank=rank, block=model.image_encoder.patch_embed.proj,
-                    dim=model.image_encoder.patch_embed.proj.out_channels
+                    rank=rank, block=model.image_encoder.patch_embed
                 )
             )
 
@@ -289,4 +276,4 @@ def __init__(
         self.sam = model
 
     def forward(self, batched_input, multimask_output):
-        return self.sam(batched_input, multimask_output)
+        return self.sam(batched_input, multimask_output)

micro_sam/training/training.py

@@ -238,7 +238,7 @@ def train_sam(
         # The number of trainable parameters for the provided model is 4.06456 (~4.06M)
 
         # peft: ssf
-        #
+        # The number of trainable parameters for the provided model is 4.267312 (~4.27M)
 
         breakpoint()