Model bileşenlerini özelleştirme

Bir modeli özelleştirmenin bir başka yolu, tamamen yeni bir model yazmak yerine bileşenlerini değiştirmektir; bu sayede modeli belirli kullanım senaryonuza göre şekillendirebilirsiniz. Örneğin, yeni katmanlar ekleyebilir veya bir mimarinin dikkat mekanizmasını optimize edebilirsiniz. Özelleştirmeler doğrudan bir Transformers modeline uygulanır, böylece Trainer, PreTrainedModel ve PEFT kütüphanesi gibi özellikleri kullanmaya devam edebilirsiniz.

Bu kılavuz, bir modelin dikkat mekanizmasını nasıl özelleştireceğinizi ve ona Low-Rank Adaptation (LoRA) uygulayacağınızı gösterecektir.

clear_import_cache yardımcı programı, model kodunu iteratif olarak değiştirip geliştirirken çok faydalıdır. Tüm önbelleğe alınmış Transformers modüllerini kaldırır ve ortamınızı sürekli yeniden başlatmadan Python'ın değiştirilen kodu yeniden yüklemesine olanak tanır.
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")

Dikkat sınıfı

Segment Anything bir görüntü segmentasyon modelidir ve dikkat mekanizmalarında sorgu-anahtar-değer (qkv) projeksiyonunu birleştirir. Eğitilebilir parametre sayısını ve hesaplama yükünü azaltmak için qkv projeksiyonuna LoRA uygulayabilirsiniz. Bunun için, LoRA ile q ve v bileşenlerini ayrı ayrı hedefleyebilmek üzere qkv projeksiyonunu bölmeniz gerekir.

SamVisionAttentionSplit adlı özel bir dikkat sınıfı oluşturmak için orijinal SamVisionAttention sınıfından türetme yapın. __init__ içinde birleşik qkv katmanını silin ve q, k ile v için ayrı ayrı doğrusal katmanlar oluşturun.

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 işlevi, model yüklenirken önceden eğitilmiş qkv ağırlıklarını ayrı q, k ve v ağırlıklarına böler ve böylece herhangi bir önceden eğitilmiş modelle uyumluluğu sağlar.

    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 geçişinde q, k ve v ayrı ayrı hesaplanırken dikkat mekanizmasının geri kalanı aynı kalır.

    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

Özel SamVisionAttentionSplit sınıfını, orijinal modelin SamVisionAttention modülüne atayarak bu modülü değiştirin. Modeldeki SamVisionAttention örneklerinin tümü, bölünmüş dikkat (split attention) sürümüyle değiştirilir.

q, k ve v katmanlarının rastgele başlatılmış olarak kalmasını önlemek için, checkpoint yüklenirken çalışacak olan yükleyici kancasının (load hook) devreye girebilmesi adına, from_pretrained() çağrısından önce sınıfı değiştirin. Zaten yüklenmiş bir modelde modülleri değiştirmek, yeni katmanların rastgele başlatılmış halde kalmasına neden olur.

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

Ayrı q, k ve v projeksiyonlarıyla, LoRA'yı q ve v üzerine uygulayın.

Bir LoraConfig oluşturun ve rank r, lora_alpha, lora_dropout, task_type ile en önemlisi, hedeflenecek modülleri belirtin.

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"
)

Modeli ve LoraConfig nesnesini get_peft_model fonksiyonuna iletin.

model = get_peft_model(model, config)

Eğittiğiniz parametre sayısını toplam parametre sayısına göre görmek için print_trainable_parameters fonksiyonunu çağırın.

model.print_trainable_parameters()
"trainable params: 589,824 || all params: 94,274,096 || trainable%: 0.6256"

Update on GitHub