micro_sam.models.peft_sam

  1import math
  2from typing import List, Union, Optional
  3
  4import torch
  5import torch.nn as nn
  6
  7from segment_anything.modeling import Sam
  8
  9try:
 10    import bitsandbytes as bnb
 11    _have_bnb = True
 12except ImportError:
 13    _have_bnb = False
 14
 15
 16class LoRASurgery(nn.Module):
 17    """Operates on the linear layers (attention and/or other feed forward) for performing low-rank adaptation.
 18
 19    (Inspired from: https://github.com/JamesQFreeman/Sam_LoRA/)
 20
 21    In SAM, it is implemented as:
 22    ```python
 23    self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
 24    B, N, C = x.shape
 25    qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
 26    q, k, v = qkv.unbind(0)
 27    ```
 28
 29    Args:
 30        rank: The rank of the decomposition matrices for updating weights in each attention layer.
 31        block: The chosen attention blocks for implementing LoRA.
 32        update_matrices: Which specific matrices to update in the attention layer. Choice of "q", "k", "v", "mlp".
 33    """
 34    def __init__(self, rank: int, block: nn.Module, update_matrices: List[str] = ["q", "v"]):
 35        super().__init__()
 36        # Check whether all values for "update_matrices" are as expected.
 37        if set(update_matrices) - set(["q", "k", "v", "mlp"]):
 38            raise ValueError(f"Some of the expected keys for updating matrics in '{update_matrices}' are not expected.")
 39
 40        self.block = block
 41        block.attn.qkv = AttentionLoRA(rank=rank, block=block.attn.qkv, update_matrices=update_matrices)
 42
 43        if "mlp" in update_matrices:
 44            block.mlp = MLPLoRA(rank=rank, mlp_layer=block.mlp)
 45
 46    def forward(self, x):
 47        return x
 48
 49
 50class AttentionLoRA(nn.Module):
 51    """Operates on the attention layers only for performing low-rank adaptation.
 52
 53    Args:
 54        rank: The rank of the decomposition matrices for updating weights in each attention layer.
 55        block: The chosen attention blocks for implementing LoRA.
 56        update_matrices: Which specific matrices to update in the attention layer. Choice of "q", "k", "v".
 57    """
 58
 59    def __init__(self, rank: int, block: nn.Module, update_matrices: List[str] = ["q", "v"]):
 60        super().__init__()
 61        self.qkv_proj = block
 62        self.dim = self.qkv_proj.in_features
 63        self.alpha = 1  # From our experiments, 'alpha' as 1 gives the best performance.
 64        self.rank = rank
 65
 66        # By default, we follow LoRA's recommended setup, i.e. update the "q" and "v" matrices.
 67        if "q" in update_matrices:
 68            self.w_a_linear_q = nn.Linear(self.dim, self.rank, bias=False)
 69            self.w_b_linear_q = nn.Linear(self.rank, self.dim, bias=False)
 70
 71        if "v" in update_matrices:
 72            self.w_a_linear_v = nn.Linear(self.dim, self.rank, bias=False)
 73            self.w_b_linear_v = nn.Linear(self.rank, self.dim, bias=False)
 74
 75        if "k" in update_matrices:
 76            self.w_a_linear_k = nn.Linear(self.dim, self.rank, bias=False)
 77            self.w_b_linear_k = nn.Linear(self.rank, self.dim, bias=False)
 78
 79        self.reset_parameters()
 80
 81        block = self
 82
 83    def reset_parameters(self):
 84        if hasattr(self, "w_a_linear_q"):
 85            nn.init.kaiming_uniform_(self.w_a_linear_q.weight, a=math.sqrt(5))
 86            nn.init.zeros_(self.w_b_linear_q.weight)
 87
 88        if hasattr(self, "w_a_linear_v"):
 89            nn.init.kaiming_uniform_(self.w_a_linear_v.weight, a=math.sqrt(5))
 90            nn.init.zeros_(self.w_b_linear_v.weight)
 91
 92        if hasattr(self, "w_a_linear_k"):
 93            nn.init.kaiming_uniform_(self.w_a_linear_k.weight, a=math.sqrt(5))
 94            nn.init.zeros_(self.w_b_linear_k.weight)
 95
 96    def forward(self, x):
 97        qkv = self.qkv_proj(x)  # B, N, N, 3 * org_C
 98
 99        new_q = self.alpha * self.w_b_linear_q(self.w_a_linear_q(x)) if hasattr(self, "w_a_linear_q") else 0
100        new_v = self.alpha * self.w_b_linear_v(self.w_a_linear_v(x)) if hasattr(self, "w_a_linear_v") else 0
101        new_k = self.alpha * self.w_b_linear_k(self.w_a_linear_k(x)) if hasattr(self, "w_a_linear_k") else 0
102        qkv = torch.cat(
103            [
104                qkv[:, :, :, :self.dim] + new_q,  # replacing new q values.
105                qkv[:, :, :, self.dim:-self.dim] + new_k,  # replacing new k values.
106                qkv[:, :, :, -self.dim:] + new_v  # replacing new v values.
107            ], dim=-1
108        )
109
110        return qkv
111
112
113class MLPLoRA(nn.Module):
114    """Operates on the feed forward layers for performing low-rank adaptation.
115
116    Args:
117        rank: The rank of the decomposition matrices for updating weights in each attention layer.
118        mlp_layer: The chosen MLP layer for implementing LoRA.
119    """
120
121    def __init__(self, rank: int, mlp_layer: nn.Module):
122        super().__init__()
123
124        self.mlp_layer = mlp_layer
125        self.rank = rank
126        self.w_a_linear_1 = nn.Linear(mlp_layer.lin1.in_features, rank, bias=False)
127        self.w_b_linear_1 = nn.Linear(rank, mlp_layer.lin1.out_features, bias=False)
128        self.w_a_linear_2 = nn.Linear(mlp_layer.lin2.in_features, rank, bias=False)
129        self.w_b_linear_2 = nn.Linear(rank, mlp_layer.lin2.out_features, bias=False)
130        self.activation = mlp_layer.act
131
132        self.reset_parameters()
133
134        mlp_layer = self
135
136    def reset_parameters(self):
137        nn.init.kaiming_uniform_(self.w_a_linear_1.weight, a=math.sqrt(5))
138        nn.init.kaiming_uniform_(self.w_a_linear_2.weight, a=math.sqrt(5))
139        nn.init.zeros_(self.w_b_linear_1.weight)
140        nn.init.zeros_(self.w_b_linear_2.weight)
141
142    def forward(self, x):
143        x = self.mlp_layer.lin1(x) + self.w_b_linear_1(self.w_a_linear_1(x))
144        x = self.activation(x)
145        x = self.mlp_layer.lin2(x) + self.w_b_linear_2(self.w_a_linear_2(x))
146        return x
147
148
149class FacTSurgery(nn.Module):
150    """Operates on the attention layers for performing factorized attention.
151
152    (Inspired from: https://github.com/cchen-cc/MA-SAM/blob/main/MA-SAM/sam_fact_tt_image_encoder.py)
153
154    Args:
155        rank: The rank of the decomposition matrices for updating weights in each attention layer.
156        block: The chosen attention blocks for implementing fact.
157        dropout: The dropout rate for the factorized attention.
158    """
159    def __init__(
160        self,
161        rank: int,
162        block: nn.Module,
163        dropout: Optional[float] = 0.1,
164    ):
165        super().__init__()
166        self.qkv_proj = block.attn.qkv
167        self.dim = self.qkv_proj.in_features
168
169        self.q_FacTs = nn.Linear(rank, rank, bias=False)
170        self.v_FacTs = nn.Linear(rank, rank, bias=False)
171
172        self.dropout = dropout
173        if self.dropout is not None:
174            self.dp_q = nn.Dropout(self.dropout)
175            self.dp_v = nn.Dropout(self.dropout)
176
177        self.FacTu = nn.Linear(self.dim, rank, bias=False)
178        self.FacTv = nn.Linear(rank, self.dim, bias=False)
179
180        block.attn.qkv = self
181
182    def forward(self, x):
183        qkv = self.qkv_proj(x)
184
185        new_q = self.q_FacTs(self.FacTu(x))
186        new_v = self.v_FacTs(self.FacTu(x))
187
188        if self.dropout is not None:
189            new_q = self.dp_q(new_q)
190            new_v = self.dp_v(new_v)
191
192        new_q = self.FacTv(new_q)
193        new_v = self.FacTv(new_v)
194
195        # NOTE : Scaling Factor is set to 1 as it can be tuned via the learning rate.
196        qkv = torch.cat(
197            [
198                qkv[:, :, :, :self.dim] + new_q,  # replacing new q values
199                qkv[:, :, :, self.dim:-self.dim],  # leaving the middle part as identical
200                qkv[:, :, :, -self.dim:] + new_v  # replacing new v values
201            ], dim=-1
202        )
203
204        return qkv
205
206
207class ScaleShiftLayer(nn.Module):
208    def __init__(self, layer, dim):
209        super().__init__()
210        self.layer = layer
211        self.scale = nn.Parameter(torch.normal(mean=1.0, std=0.2, size=(dim,)))
212        self.shift = nn.Parameter(torch.normal(mean=0.0, std=0.2, size=(dim,)))
213        layer = self
214
215    def forward(self, x):
216        x = self.layer(x)
217        assert self.scale.shape == self.shift.shape
218        if x.shape[-1] == self.scale.shape[0]:
219            return x * self.scale + self.shift
220        elif x.shape[1] == self.scale.shape[0]:
221            return x * self.scale.view(1, -1, 1, 1) + self.shift.view(1, -1, 1, 1)
222        else:
223            raise ValueError('Input tensors do not match the shape of the scale factors.')
224
225
226class SSFSurgery(nn.Module):
227    """Operates on all layers in the transformer block for adding learnable scale and shift parameters.
228
229    Args:
230        rank: This parameter is not used in `SSFSurgery`. This is kept here for consistency.
231        block: The chosen attention blocks for implementing ssf.
232    """
233    def __init__(self, rank: int, block: nn.Module):
234        super().__init__()
235        self.block = block
236
237        # If we get a transformer block (w. multiple sub-layers), we perform surgery on each layer.
238        if hasattr(block, "attn"):  # the minimum assumption is to verify the attention layers.
239            block.attn.qkv = ScaleShiftLayer(block.attn.qkv, block.attn.qkv.in_features*3)
240            block.attn.proj = ScaleShiftLayer(block.attn.proj, block.attn.proj.in_features)
241            block.mlp.lin1 = ScaleShiftLayer(block.mlp.lin1, block.mlp.lin1.out_features)
242            block.mlp.lin2 = ScaleShiftLayer(block.mlp.lin2, block.mlp.lin2.out_features)
243            block.norm1 = ScaleShiftLayer(block.norm1, block.norm1.normalized_shape[0])
244            block.norm2 = ScaleShiftLayer(block.norm2, block.norm2.normalized_shape[0])
245
246        # If we get the embedding block, add one ScaleShiftLayer
247        elif hasattr(block, "patch_embed"):
248            block.proj = ScaleShiftLayer(block.proj, block.proj.out_channels)
249
250    def forward(self, x):
251        return x
252
253
254class SelectiveSurgery(nn.Module):
255    """Base class for selectively allowing gradient updates for certain parameters.
256    """
257    def __init__(self, block: nn.Module):
258        super().__init__()
259        self.block = block
260
261    def allow_gradient_update_for_parameters(
262        self,
263        prefix: Optional[List[str]] = None,
264        suffix: Optional[List[str]] = None,
265        infix: Optional[List[str]] = None,
266    ):
267        """This function decides the parameter attributes to match for allowing gradient updates.
268
269        Args:
270            prefix: Matches the part of parameter name in front.
271            suffix: Matches the part of parameter name at the end.
272            infix: Matches parts of parameter name occuring in between.
273        """
274        for k, v in self.block.named_parameters():
275            if prefix is not None and k.startswith(tuple(prefix)):
276                v.requires_grad = True
277
278            if suffix is not None and k.endswith(tuple(suffix)):
279                v.requires_grad = True
280
281            if infix is not None:
282                for per_infix in infix:
283                    if k.find(per_infix) != -1:
284                        v.requires_grad = True
285
286    def forward(self, x):
287        return x
288
289
290class AdaptFormer(nn.Module):
291    """Adds AdaptFormer Module in place of the MLP Layers
292
293    Args:
294        rank: The rank is not used in this class but kept here for consistency.
295        block: The chosen encoder block for implementing AdaptFormer.
296        alpha: A parameters that scales the Adapter path. Can be either learnable or some fixed value.
297        dropout: The dropout rate for the dropout layer between down and up projection layer.
298        projection_size: The size of the projection layer.
299    """
300    def __init__(
301        self,
302        rank: int,
303        block: nn.Module,
304        alpha: Optional[Union[str, float]] = "learnable_scalar",  # Stable choice from our preliminary exp.
305        dropout: Optional[float] = None,  # Does not have an obvious advantage.
306        projection_size: int = 64,  # Stable choice from our preliminary exp.
307    ):
308        super().__init__()
309
310        self.mlp_proj = block.mlp
311        self.n_embd = block.mlp.lin1.in_features
312
313        if alpha == 'learnable_scalar':
314            self.alpha = nn.Parameter(torch.ones(1))
315        else:
316            self.alpha = alpha
317
318        self.projection_size = projection_size
319        self.dropout = dropout
320
321        self.down_proj = nn.Linear(self.n_embd, self.projection_size)
322        self.non_linear_func = nn.ReLU()
323        self.up_proj = nn.Linear(self.projection_size, self.n_embd)
324
325        block.mlp = self
326
327        if self.dropout is not None:
328            self.dropout_layer = nn.Dropout(self.dropout)
329
330        nn.init.kaiming_uniform_(self.down_proj.weight, a=math.sqrt(5))
331        nn.init.zeros_(self.up_proj.weight)
332        nn.init.zeros_(self.down_proj.bias)
333        nn.init.zeros_(self.up_proj.bias)
334
335    def forward(self, x):
336        residual = x
337        mlp_output = self.mlp_proj(x)
338
339        down = self.down_proj(x)
340        down = self.non_linear_func(down)
341
342        if self.dropout is not None:
343            down = self.dropout_layer(down)
344
345        up = self.up_proj(down)
346        up = up * self.alpha
347        output = up + residual + mlp_output
348
349        return output
350
351
352class AttentionSurgery(SelectiveSurgery):
353    """Child class for allowing gradient updates for parameters in attention layers."""
354
355    def __init__(self, block: nn.Module):
356        super().__init__(block=block)
357        # Allow gradient updates for the attention layers in the image encoder.
358        self.allow_gradient_update_for_parameters(prefix=["attn"])
359
360
361class BiasSurgery(SelectiveSurgery):
362    """Child class for allowing gradient updates for bias parameters."""
363
364    def __init__(self, block: nn.Module):
365        super().__init__(block=block)
366        # Allow gradient updates for the bias parameters in the image encoder.
367        self.allow_gradient_update_for_parameters(suffix=["bias"])
368
369
370class LayerNormSurgery(SelectiveSurgery):
371    """Child class for allowing gradient updates in normalization layers."""
372
373    def __init__(self, block: nn.Module):
374        super().__init__(block=block)
375        # Allow gradient updates for the LayerNorm parameters in the image encoder.
376        self.allow_gradient_update_for_parameters(infix=["norm1", "norm2"])
377
378
379class ClassicalSurgery(SelectiveSurgery):
380    """Child class for freezing specific blocks."""
381
382    def __init__(self, block: nn.Module):
383        super().__init__(block=block)
384        self.block = block
385
386        for k, v in self.block.named_parameters():
387            v.requires_grad = True
388
389    def forward(self, x):
390        return x
391
392
393class PEFT_Sam(nn.Module):
394    """Wraps the Segment Anything model's image encoder to different parameter efficient finetuning methods.
395
396    Args:
397        model: The Segment Anything model.
398        rank: The rank for low-rank adaptation.
399        peft_module: Wrapper to operate on the image encoder blocks for the PEFT method.
400        attention_layers_to_update: Which specific layers we apply PEFT methods to.
401            For reference, the total number of blocks for 'vit_b' is 12, for 'vit_l' is 24 and for 'vit_h' is 32.
402        quantize: Whether to quantize the model for lower precision training.
403        module_kwargs: The additional arguments for the respective PEFT modules.
404    """
405
406    def __init__(
407        self,
408        model: Sam,
409        rank: Optional[int] = None,
410        peft_module: nn.Module = LoRASurgery,
411        attention_layers_to_update: Optional[List[int]] = None,
412        quantize: bool = False,
413        **module_kwargs
414    ):
415        super().__init__()
416
417        if issubclass(peft_module, Union[LoRASurgery, FacTSurgery]) and (not rank or rank <= 0):
418            raise RuntimeError("The chosen PEFT method cannot run without a valid rank choice.")
419
420        assert issubclass(peft_module, Union[LoRASurgery, FacTSurgery, SelectiveSurgery, SSFSurgery, AdaptFormer]), (
421            "Invalid PEFT module"
422        )
423        if attention_layers_to_update:
424            self.peft_layers = attention_layers_to_update
425        else:   # Applies PEFT to the image encoder by default
426            self.peft_layers = list(range(len(model.image_encoder.blocks)))
427
428        self.peft_module = peft_module
429        self.peft_blocks = []
430
431        # Whether to quantize the linear layers to 4 bit precision.
432        # NOTE: This is currently supported for CUDA-supported devices only.
433        if quantize:
434            if not _have_bnb:
435                raise ModuleNotFoundError("Please install 'bitsandbytes'.")
436
437            for name, module in model.image_encoder.named_modules():
438                if isinstance(module, torch.nn.Linear):
439                    *parent_path, layer_name = name.split(".")
440                    parent_module = model.image_encoder
441
442                    for sub_module in parent_path:
443                        parent_module = getattr(parent_module, sub_module)
444
445                    # Create the new Linear4bit layer
446                    linear_q = bnb.nn.Linear4bit(
447                        module.in_features,
448                        module.out_features,
449                        bias=False if module.bias is None else True,
450                    )
451                    # Assign weights and bias to the new layer
452                    new_weight = bnb.nn.Params4bit(
453                        data=module.weight,
454                        requires_grad=False,
455                    )
456                    linear_q.weight = new_weight
457                    if module.bias is not None:
458                        linear_q.bias = torch.nn.Parameter(module.bias)
459
460                    # Replace the original linear layer with the quantized one
461                    setattr(parent_module, layer_name, linear_q)
462
463        # Let's freeze all the pretrained image encoder layers first
464        for param in model.image_encoder.parameters():
465            param.requires_grad = False
466
467        # Add scale and shift parameters to the patch embedding layers.
468        if issubclass(self.peft_module, SSFSurgery):
469            self.peft_blocks.append(self.peft_module(rank=rank, block=model.image_encoder.patch_embed))
470
471        # If specified, the attention layers to update should match the available blocks.
472        if attention_layers_to_update and (
473            set(attention_layers_to_update) - set(list(range(len(model.image_encoder.blocks))))
474        ):
475            raise ValueError("The chosen layer(s) to apply PEFT method is not a valid transformer block id.")
476
477        for t_layer_i, blk in enumerate(model.image_encoder.blocks):
478
479            # If we only want specific layers with PEFT instead of all
480            if t_layer_i not in self.peft_layers:
481                continue
482
483            if issubclass(self.peft_module, SelectiveSurgery):
484                self.peft_blocks.append(self.peft_module(block=blk))
485            else:
486                self.peft_blocks.append(self.peft_module(rank=rank, block=blk, **module_kwargs))
487
488        self.peft_blocks = nn.ModuleList(self.peft_blocks)
489        self.sam = model
490
491    def forward(self, batched_input, multimask_output):
492        return self.sam(batched_input, multimask_output)