모델 구성 요소 사용자 정의

모델을 커스터마이징하는 또 다른 방법은 모델을 처음부터 새로 작성하는 대신, 그 구성 요소를 수정하는 것입니다. 이를 통해 특정 사용 사례에 맞게 모델을 조정할 수 있습니다. 예를 들어, 새로운 레이어를 추가하거나 아키텍처의 어텐션 메커니즘을 최적화할 수 있습니다. 이러한 커스터마이징은 Trainer, PreTrainedModel, PEFT 라이브러리 등의 기능을 계속 사용할 수 있도록 Transformers 모델에 직접 적용됩니다.

이 가이드에서는 모델의 attention mechanism을 커스터마이징하여 Low-Rank Adaptation (LoRA)를 적용하는 방법을 보여줍니다.

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는 이미지 분할 모델로, 어텐션 메커니즘에서 쿼리-키-밸류(qkv) 투영을 결합합니다. 학습 가능한 매개변수 수와 계산 오버헤드를 줄이려면 qkv 투영에 LoRA를 적용할 수 있습니다. 이를 위해 qkv 투영을 분리하여 q와 v에 LoRA를 개별적으로 적용할 수 있도록 해야 합니다.

SamVisionAttention 클래스를 서브클래스로 만들어 커스텀 어텐션 클래스 SamVisionAttentionSplit를 생성합니다. __init__에서 결합된 qkv를 삭제하고 q, k, v 각각에 대해 별도의 선형 레이어를 생성합니다.

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]

forward 패스에서 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 인스턴스가 분리된 어텐션 버전으로 교체됩니다.

q, k, v 레이어가 무작위로 초기화되는 것을 방지하려면, from_pretrained()를 호출하기 전에 클래스를 교체해야 합니다. 로드 훅은 체크포인트가 로드되는 동안에만 실행되므로, 이미 로드된 모델에서 모듈을 교체하면 새로운 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를 생성하고 랭크 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"

Update on GitHub