Настройка компонентов модели

Ещё один способ настроить модель — изменить её компоненты, а не писать новую модель с нуля, что позволяет адаптировать модель под конкретную задачу. Например, можно добавить новые слои или оптимизировать механизм внимания в архитектуре. Настройка применяется непосредственно к модели Transformers, поэтому можно продолжать использовать такие функции, как Trainer, PreTrainedModel и библиотека PEFT.

В этом руководстве показано, как настроить механизм внимания модели для применения к нему 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")

Класс внимания

Segment Anything — это модель сегментации изображений, которая объединяет проекцию query-key-value (qkv) в своих механизмах внимания. Чтобы сократить количество обучаемых параметров и вычислительные затраты, вы можете применить LoRA к проекции qkv. Для этого необходимо разделить проекцию qkv, чтобы можно было отдельно нацелить LoRA на q и v.

Создайте собственный класс внимания SamVisionAttentionSplit, наследуясь от исходного класса SamVisionAttention. В __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 вычисляются отдельно, при этом остальная часть механизма внимания остаётся без изменений.

    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 в модели заменяются версией с разделённым вниманием.

Замените класс до вызова 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 применяйте LoRA к q и v.

Создайте 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