모델 구성요소 커스터마이징
모델 구성요소 커스터마이징 (Customizing model components)
모델을 커스터마이즈하는 또 다른 방법은 아예 새 모델을 작성하는 대신 그 구성요소를 수정하는 거예요. 이를 통해 모델을 특정 사용 사례에 맞출 수 있어요. 예를 들어 새 레이어를 추가하거나 아키텍처의 attention 메커니즘을 최적화할 수 있어요. 커스터마이즈는 Transformers 모델에 직접 적용되므로 Trainer, PreTrainedModel, PEFT 라이브러리 같은 기능을 계속 사용할 수 있어요.
이 가이드에서는 Low-Rank Adaptation (LoRA)을 적용하기 위해 모델의 attention 메커니즘을 커스터마이즈하는 방법을 보여줄게요.
[!TIP] clear_import_cache 유틸리티는 모델 코드를 반복적으로 수정하고 개발할 때 매우 유용합니다. 캐시된 Transformers 모듈을 모두 제거하고 환경을 계속 재시작하지 않아도 Python이 수정된 코드를 다시 로드하게 해줍니다.
from transformers import AutoModel from transformers.utils.import_utils import clear_import_cache model = AutoModel.from_pretrained("bert-base-uncased") # modifications to model code # clear cache to reload modified code clear_import_cache() # re-import to use updated code model = AutoModel.from_pretrained("bert-base-uncased")
출처: 문서
본문
Attention 클래스
Segment Anything은 이미지 분할 모델로, attention 메커니즘에서 query-key-value (qkv) 프로젝션을 결합해요. 학습 가능한 파라미터 수와 계산 오버헤드를 줄이기 위해 qkv 프로젝션에 LoRA를 적용할 수 있어요. 이를 위해서는 qkv 프로젝션을 분할해서 q와 v를 LoRA로 개별적으로 타겟팅할 수 있어야 해요.
- 원래
SamVisionAttention클래스를 서브클래싱한 커스텀 attention 클래스SamVisionAttentionSplit을 만들어요.__init__에서 결합된qkv를 삭제하고q,k,v각각에 별도의 linear 레이어를 만들어요.
import torch
import torch.nn as nn
from transformers.models.sam.modeling_sam import SamVisionAttention
class SamVisionAttentionSplit(SamVisionAttention, nn.Module):
def __init__(self, config, window_size):
super().__init__(config, window_size)
# remove combined qkv
del self.qkv
# separate q, k, v projections
self.q = nn.Linear(config.hidden_size, config.hidden_size, bias=config.qkv_bias)
self.k = nn.Linear(config.hidden_size, config.hidden_size, bias=config.qkv_bias)
self.v = nn.Linear(config.hidden_size, config.hidden_size, bias=config.qkv_bias)
self._register_load_state_dict_pre_hook(self.split_q_k_v_load_hook)
split_q_k_v_load_hook함수는 모델을 로드할 때 사전 학습된qkv가중치를 별도의q,k,v가중치로 분할해서 어떤 사전 학습된 모델과도 호환되게 해요.
def split_q_k_v_load_hook(self, state_dict, prefix, *args):
keys_to_delete = []
for key in list(state_dict.keys()):
if "qkv." in key:
# split q, k, v from the combined projection
q, k, v = state_dict[key].chunk(3, dim=0)
# replace with individual q, k, v projections
state_dict[key.replace("qkv.", "q.")] = q
state_dict[key.replace("qkv.", "k.")] = k
state_dict[key.replace("qkv.", "v.")] = v
# mark the old qkv key for deletion
keys_to_delete.append(key)
# remove old qkv keys
for key in keys_to_delete:
del state_dict[key]
forwardpass에서q,k,v는 별도로 계산되며 나머지 attention 메커니즘은 동일하게 유지돼요.
def forward(self, hidden_states: torch.Tensor, output_attentions=False) -> torch.Tensor:
batch_size, height, width, _ = hidden_states.shape
qkv_shapes = (batch_size * self.num_attention_heads, height * width, -1)
query = self.q(hidden_states).reshape((batch_size, height * width,self.num_attention_heads, -1)).permute(0,2,1,3).reshape(qkv_shapes)
key = self.k(hidden_states).reshape((batch_size, height * width,self.num_attention_heads, -1)).permute(0,2,1,3).reshape(qkv_shapes)
value = self.v(hidden_states).reshape((batch_size, height * width,self.num_attention_heads, -1)).permute(0,2,1,3).reshape(qkv_shapes)
attn_weights = (query * self.scale) @ key.transpose(-2, -1)
attn_weights = torch.nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query.dtype)
attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
attn_output = (attn_probs @ value).reshape(batch_size, self.num_attention_heads, height, width, -1)
attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1)
attn_output = self.proj(attn_output)
if output_attentions:
outputs = (attn_output, attn_weights)
else:
outputs = (attn_output, None)
return outputs
커스텀 SamVisionAttentionSplit 클래스를 원래 모델의 SamVisionAttention 모듈에 할당해서 교체해요. 모델에 있는 모든 SamVisionAttention 인스턴스가 분할 attention 버전으로 교체돼요.
from_pretrained()을 호출하기 전에 클래스를 교체하세요. load hook은 체크포인트를 로드하는 동안에만 실행되므로, 이미 로드된 모델의 모듈을 교체하면 새 q, k, v 레이어가 무작위로 초기화된 채로 남게 돼요.
from transformers import SamModel
from transformers.models.sam import modeling_sam
# replace the attention class the vision layers instantiate
modeling_sam.SAM_VISION_ATTENTION_CLASSES["eager"] = SamVisionAttentionSplit
# load the pretrained SAM model, the hook splits the qkv weights during loading
model = SamModel.from_pretrained("facebook/sam-vit-base", attn_implementation="eager")
LoRA
q, k, v 프로젝션이 분리되었으니 q와 v에 LoRA를 적용해요.
LoraConfig를 만들고 rank r, lora_alpha, lora_dropout, task_type, 그리고 가장 중요한 타겟팅할 모듈을 지정해요.
from peft import LoraConfig, get_peft_model
config = LoraConfig(
r=16,
lora_alpha=32,
# apply LoRA to q and v
target_modules=["q", "v"],
lora_dropout=0.1,
task_type="FEATURE_EXTRACTION"
)
모델과 LoraConfig를 get_peft_model에 전달해서 모델에 LoRA를 적용해요.
model = get_peft_model(model, config)
print_trainable_parameters를 호출해서 결과적으로 학습하는 파라미터 수를 총 파라미터 수와 비교해 확인해요.
model.print_trainable_parameters()
"trainable params: 589,824 || all params: 94,274,096 || trainable%: 0.6256"