| from segment_anything import build_sam, SamPredictor
|
| from segment_anything import sam_model_registry
|
|
|
| import math
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from torch import Tensor
|
| from torch.nn.parameter import Parameter
|
| from segment_anything.modeling import Sam
|
| from safetensors import safe_open
|
| from safetensors.torch import save_file
|
|
|
| from icecream import ic
|
|
|
|
|
| class _LoRA_qkv(nn.Module):
|
| """In Sam it is implemented as
|
| self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
| B, N, C = x.shape
|
| qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
|
| q, k, v = qkv.unbind(0)
|
| """
|
|
|
| def __init__(
|
| self,
|
| qkv: nn.Module,
|
| linear_a_q: nn.Module,
|
| linear_b_q: nn.Module,
|
| linear_a_v: nn.Module,
|
| linear_b_v: nn.Module,
|
| ):
|
| super().__init__()
|
| self.qkv = qkv
|
| self.linear_a_q = linear_a_q
|
| self.linear_b_q = linear_b_q
|
| self.linear_a_v = linear_a_v
|
| self.linear_b_v = linear_b_v
|
| self.dim = qkv.in_features
|
| self.w_identity = torch.eye(qkv.in_features)
|
|
|
| def forward(self, x):
|
| qkv = self.qkv(x)
|
| new_q = self.linear_b_q(self.linear_a_q(x))
|
| new_v = self.linear_b_v(self.linear_a_v(x))
|
| qkv[:, :, :, : self.dim] += new_q
|
| qkv[:, :, :, -self.dim:] += new_v
|
| return qkv
|
|
|
|
|
| class LoRA_Sam(nn.Module):
|
| """Applies low-rank adaptation to a Sam model's image encoder.
|
|
|
| Args:
|
| sam_model: a vision transformer model, see base_vit.py
|
| r: rank of LoRA
|
| num_classes: how many classes the model output, default to the vit model
|
| lora_layer: which layer we apply LoRA.
|
|
|
| Examples::
|
| >>> model = ViT('B_16_imagenet1k')
|
| >>> lora_model = LoRA_ViT(model, r=4)
|
| >>> preds = lora_model(img)
|
| >>> print(preds.shape)
|
| torch.Size([1, 1000])
|
| """
|
|
|
| def __init__(self, sam_model: Sam, r: int, lora_layer=None):
|
| super(LoRA_Sam, self).__init__()
|
|
|
| assert r > 0
|
|
|
|
|
| if lora_layer:
|
| self.lora_layer = lora_layer
|
| else:
|
| self.lora_layer = list(
|
| range(len(sam_model.image_encoder.blocks)))
|
|
|
| self.w_As = []
|
| self.w_Bs = []
|
|
|
|
|
| for param in sam_model.image_encoder.parameters():
|
| param.requires_grad = False
|
|
|
|
|
| for t_layer_i, blk in enumerate(sam_model.image_encoder.blocks):
|
|
|
| if t_layer_i not in self.lora_layer:
|
| continue
|
| w_qkv_linear = blk.attn.qkv
|
| self.dim = w_qkv_linear.in_features
|
| w_a_linear_q = nn.Linear(self.dim, r, bias=False)
|
| w_b_linear_q = nn.Linear(r, self.dim, bias=False)
|
| w_a_linear_v = nn.Linear(self.dim, r, bias=False)
|
| w_b_linear_v = nn.Linear(r, self.dim, bias=False)
|
| self.w_As.append(w_a_linear_q)
|
| self.w_Bs.append(w_b_linear_q)
|
| self.w_As.append(w_a_linear_v)
|
| self.w_Bs.append(w_b_linear_v)
|
| blk.attn.qkv = _LoRA_qkv(
|
| w_qkv_linear,
|
| w_a_linear_q,
|
| w_b_linear_q,
|
| w_a_linear_v,
|
| w_b_linear_v,
|
| )
|
| self.reset_parameters()
|
| self.sam = sam_model
|
|
|
| def save_lora_parameters(self, filename: str) -> None:
|
| r"""Only safetensors is supported now.
|
|
|
| pip install safetensor if you do not have one installed yet.
|
|
|
| save both lora and fc parameters.
|
| """
|
|
|
| assert filename.endswith(".pt") or filename.endswith('.pth')
|
|
|
| num_layer = len(self.w_As)
|
| a_tensors = {f"w_a_{i:03d}": self.w_As[i].weight for i in range(num_layer)}
|
| b_tensors = {f"w_b_{i:03d}": self.w_Bs[i].weight for i in range(num_layer)}
|
| prompt_encoder_tensors = {}
|
| mask_decoder_tensors = {}
|
|
|
|
|
| if isinstance(self.sam, torch.nn.DataParallel) or isinstance(self.sam, torch.nn.parallel.DistributedDataParallel):
|
| state_dict = self.sam.module.state_dict()
|
| else:
|
| state_dict = self.sam.state_dict()
|
| for key, value in state_dict.items():
|
| if 'prompt_encoder' in key:
|
| prompt_encoder_tensors[key] = value
|
| if 'mask_decoder' in key:
|
| mask_decoder_tensors[key] = value
|
|
|
| merged_dict = {**a_tensors, **b_tensors, **prompt_encoder_tensors, **mask_decoder_tensors}
|
| torch.save(merged_dict, filename)
|
|
|
| def load_lora_parameters(self, filename: str) -> None:
|
| r"""Only safetensors is supported now.
|
|
|
| pip install safetensor if you do not have one installed yet.\
|
|
|
| load both lora and fc parameters.
|
| """
|
|
|
| assert filename.endswith(".pt") or filename.endswith('.pth')
|
|
|
| state_dict = torch.load(filename)
|
|
|
| for i, w_A_linear in enumerate(self.w_As):
|
| saved_key = f"w_a_{i:03d}"
|
| saved_tensor = state_dict[saved_key]
|
| w_A_linear.weight = Parameter(saved_tensor)
|
|
|
| for i, w_B_linear in enumerate(self.w_Bs):
|
| saved_key = f"w_b_{i:03d}"
|
| saved_tensor = state_dict[saved_key]
|
| w_B_linear.weight = Parameter(saved_tensor)
|
|
|
| sam_dict = self.sam.state_dict()
|
| sam_keys = sam_dict.keys()
|
|
|
|
|
| prompt_encoder_keys = [k for k in sam_keys if 'prompt_encoder' in k]
|
| prompt_encoder_values = [state_dict[k] for k in prompt_encoder_keys]
|
| prompt_encoder_new_state_dict = {k: v for k, v in zip(prompt_encoder_keys, prompt_encoder_values)}
|
| sam_dict.update(prompt_encoder_new_state_dict)
|
|
|
|
|
| mask_decoder_keys = [k for k in sam_keys if 'mask_decoder' in k]
|
| mask_decoder_values = [state_dict[k] for k in mask_decoder_keys]
|
| mask_decoder_new_state_dict = {k: v for k, v in zip(mask_decoder_keys, mask_decoder_values)}
|
| sam_dict.update(mask_decoder_new_state_dict)
|
| self.sam.load_state_dict(sam_dict)
|
|
|
| def reset_parameters(self) -> None:
|
| for w_A in self.w_As:
|
| nn.init.kaiming_uniform_(w_A.weight, a=math.sqrt(5))
|
| for w_B in self.w_Bs:
|
| nn.init.zeros_(w_B.weight)
|
|
|
| def forward(self, batched_input, multimask_output, image_size):
|
| return self.sam(batched_input, multimask_output, image_size)
|
|
|
|
|
|
|
|
|
|
|
|
|
| if __name__ == "__main__":
|
| sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth")
|
| lora_sam = LoRA_Sam(sam, 4)
|
| lora_sam.sam.image_encoder(torch.rand(size=(1, 3, 1024, 1024)))
|
|
|
