Personnalisation des composants du modèle

Une autre façon de personnaliser un modèle consiste à modifier ses composants, plutôt que d’écrire un nouveau modèle de zéro, ce qui vous permet d’adapter un modèle à votre cas d’usage spécifique. Par exemple, vous pouvez ajouter de nouvelles couches ou optimiser le mécanisme d’attention d’une architecture. Les personnalisations sont appliquées directement à un modèle Transformers afin que vous puissiez continuer à utiliser des fonctionnalités telles que Trainer, PreTrainedModel et la bibliothèque PEFT.

Ce guide vous montre comment personnaliser le mécanisme d'attention d'un modèle afin d'y appliquer Low-Rank Adaptation (LoRA).

L'utilitaire clear_import_cache est très utile lorsque vous modifiez et développez itérativement le code d'un modèle. Il supprime tous les modules Transformers mis en cache et permet à Python de recharger le code modifié sans avoir à redémarrer constamment votre environnement.
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")

Classe d'attention

Segment Anything est un modèle de segmentation d’images qui combine la projection query-key-value (qkv) dans ses mécanismes d’attention. Pour réduire le nombre de paramètres entraînables et la surcharge de calcul, vous pouvez appliquer LoRA à la projection qkv. Cela nécessite de découper la projection qkv afin de pouvoir cibler séparément q et v avec LoRA.

Créez une classe d'attention personnalisée, SamVisionAttentionSplit, en sous-classant la classe originale SamVisionAttention. Dans __init__, supprimez le qkv combiné et créez une couche linéaire distincte pour q, k et 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)

La fonction split_q_k_v_load_hook sépare les poids préentraînés qkv en poids distincts q, k et v lors du chargement du modèle afin d'assurer la compatibilité avec tout modèle préentraîné.

    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]

Lors du passage forward, q, k et v sont calculés séparément, tandis que le reste du mécanisme d'attention reste identique.

    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

Affectez la classe personnalisée SamVisionAttentionSplit au module SamVisionAttention des modèles d'origine pour le remplacer. Toutes les instances de SamVisionAttention dans le modèle sont remplacées par la version d'attention scindée.

Remplacez la classe avant d’appeler from_pretrained(). Le hook de chargement ne s’exécute que pendant le chargement du point de contrôle ; remplacer les modules sur un modèle déjà chargé laisse les nouvelles couches q, k et v initialisées aléatoirement.

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

Avec des projections q, k et v distinctes, appliquez LoRA à q et v.

Créez un LoraConfig et spécifiez le rang r, lora_alpha, lora_dropout, task_type et, surtout, les modules à cibler.

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

Passez le modèle et LoraConfig à get_peft_model pour appliquer LoRA au modèle.

model = get_peft_model(model, config)

Appel de print_trainable_parameters pour afficher le nombre de paramètres entraînés par rapport au nombre total de paramètres.

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

Update on GitHub