Menyesuaikan komponen model

Cara lain untuk menyesuaikan model adalah dengan memodifikasi komponen-komponennya, alih-alih menulis model baru sepenuhnya, sehingga Anda dapat menyesuaikan model untuk kasus penggunaan spesifik Anda. Misalnya, Anda dapat menambahkan lapisan baru atau mengoptimalkan mekanisme perhatian dari sebuah arsitektur. Penyesuaian diterapkan langsung pada model Transformers sehingga Anda dapat terus menggunakan fitur-fitur seperti Trainer, PreTrainedModel, dan pustaka PEFT.

Panduan ini akan menunjukkan cara menyesuaikan mekanisme perhatian model untuk menerapkan Low-Rank Adaptation (LoRA) padanya.

Utilitas clear_import_cache sangat berguna saat Anda secara iteratif memodifikasi dan mengembangkan kode model. Utilitas ini menghapus semua modul Transformers yang di-cache dan memungkinkan Python untuk memuat ulang kode yang telah dimodifikasi tanpa perlu terus-menerus memulai ulang lingkungan Anda.
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")

Kelas Attention

Segment Anything adalah model segmentasi gambar, dan menggabungkan proyeksi query-key-value (qkv) dalam mekanisme atensinya. Untuk mengurangi jumlah parameter yang dapat dilatih dan beban komputasi, Anda dapat menerapkan LoRA pada proyeksi qkv. Hal ini memerlukan pemisahan proyeksi qkv agar Anda dapat menargetkan q dan v secara terpisah dengan LoRA.

Buat kelas perhatian kustom, SamVisionAttentionSplit, dengan mewarisi kelas SamVisionAttention asli. Di dalam __init__, hapus qkv yang digabungkan dan buat lapisan linear terpisah untuk q, k, dan 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)

Fungsi split_q_k_v_load_hook memisahkan bobot qkv yang telah dilatih sebelumnya menjadi bobot q, k, dan v yang terpisah saat memuat model untuk memastikan kompatibilitas dengan model yang telah dilatih sebelumnya.

    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]

Pada langkah forward, q, k, dan v dihitung secara terpisah, sementara mekanisme perhatian lainnya tetap sama.

    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

Tugaskan kelas SamVisionAttentionSplit kustom ke modul SamVisionAttention model asli untuk menggantikannya. Semua instans SamVisionAttention dalam model diganti dengan versi perhatian terpecah.

Ganti kelas sebelum memanggil from_pretrained(). Hook pemuatan hanya berjalan saat checkpoint dimuat, sehingga mengganti modul pada model yang sudah dimuat akan membuat lapisan q, k, dan v baru terinisialisasi secara acak.

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

Dengan proyeksi q, k, dan v yang terpisah, terapkan LoRA pada q dan v.

Buat LoraConfig dan tentukan rank r, lora_alpha, lora_dropout, task_type, dan yang paling penting, modul yang akan ditargetkan.

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

Lewatkan model dan LoraConfig ke get_peft_model untuk menerapkan LoRA ke model.

model = get_peft_model(model, config)

Panggil print_trainable_parameters untuk melihat jumlah parameter yang Anda latih sebagai hasilnya dibandingkan dengan total jumlah parameter.

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

Update on GitHub