Personalizar los componentes del modelo

Otra forma de personalizar un modelo es modificar sus componentes, en lugar de escribir un modelo nuevo por completo, lo que te permite adaptar un modelo a tu caso de uso específico. Por ejemplo, puedes agregar nuevas capas u optimizar el mecanismo de atención de una arquitectura. Las personalizaciones se aplican directamente a un modelo de Transformers para que puedas seguir utilizando funciones como Trainer, PreTrainedModel y la biblioteca PEFT.

Esta guía le mostrará cómo personalizar el mecanismo de atención de un modelo para aplicar Low-Rank Adaptation (LoRA) a él.

La utilidad clear_import_cache es muy útil cuando se modifica y desarrolla el código de un modelo de forma iterativa. Elimina todos los módulos de Transformers en caché y permite que Python recargue el código modificado sin necesidad de reiniciar constantemente el entorno.
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")

Clase de atención

Segment Anything es un modelo de segmentación de imágenes y combina la proyección query-key-value (qkv) en sus mecanismos de atención. Para reducir el número de parámetros entrenables y la sobrecarga computacional, puedes aplicar LoRA a la proyección qkv. Esto requiere dividir la proyección qkv para poder aplicar LoRA por separado a q y v.

Crea una clase de atención personalizada, SamVisionAttentionSplit, subclaseando la clase original SamVisionAttention. En __init__, elimina el qkv combinado y crea una capa lineal separada para q, k y 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 función split_q_k_v_load_hook divide los pesos preentrenados de qkv en pesos separados de q, k y v al cargar el modelo para garantizar la compatibilidad con cualquier modelo preentrenado.

    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]

En la pasada forward, q, k y v se calculan por separado, mientras que el resto del mecanismo de atención permanece igual.

    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

Asigne la clase personalizada SamVisionAttentionSplit al módulo SamVisionAttention de los modelos originales para reemplazarlo. Todas las instancias de SamVisionAttention en el modelo se reemplazan con la versión de atención dividida.

Sustituya la clase antes de llamar a from_pretrained(). El hook de carga solo se ejecuta mientras se carga el checkpoint, por lo que reemplazar los módulos en un modelo ya cargado deja las nuevas capas q, k y v inicializadas aleatoriamente.

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

Con proyecciones q, k y v separadas, aplica LoRA a q y v.

Cree una LoraConfig y especifique el rango r, lora_alpha, lora_dropout, task_type y, lo más importante, los módulos a los que se aplicará.

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

Pasa el modelo y LoraConfig a get_peft_model para aplicar LoRA al modelo.

model = get_peft_model(model, config)

Llama a print_trainable_parameters para ver el número de parámetros que estás entrenando en comparación con el número total de parámetros.

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

Update on GitHub