# モデルコンポーネントのカスタマイズ

モデルをカスタマイズするもう一つの方法は、新しいモデルをゼロから書くのではなく、既存のコンポーネントを変更することです。これにより、特定のユースケースに合わせてモデルを調整できます。例えば、新しいレイヤーを追加したり、アーキテクチャの注意機構を最適化したりできます。カスタマイズは直接 Transformers のモデルに適用されるため、[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer)、[PreTrainedModel](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel)、および [PEFT](https://huggingface.co/docs/peft/en/index) ライブラリなどの機能を引き続き利用できます。

このガイドでは、モデルの[Low-Rank Adaptation (LoRA)](https://huggingface.co/docs/peft/conceptual_guides/adapter#low-rank-adaptation-lora)を適用するために、アテンションメカニズムをカスタマイズする方法を示します。

> [!TIP]
> [clear_import_cache](https://github.com/huggingface/transformers/blob/9985d06add07a4cc691dc54a7e34f54205c04d40/src/transformers/utils/import_utils.py#L2286) ユーティリティは、モデルコードを反復的に修正・開発する際に非常に役立ちます。すべての Transformers モジュールのキャッシュを削除し、環境を頻繁に再起動することなく、Python が修正されたコードを再読み込みできるようにします。
>
> ```py
> 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")
> ```

## Attentionクラス

[Segment Anything](./model_doc/sam)は画像セグメンテーションモデルであり、そのアテンション機構においてクエリ・キー・バリュー（`qkv`）の射影を結合しています。学習可能なパラメータ数と計算オーバーヘッドを削減するために、`qkv`射影にLoRAを適用できます。これには、`qkv`射影を分割して、`q`と`v`を個別にLoRAのターゲットとして指定できるようにする必要があります。

1. `SamVisionAttention` クラスをサブクラス化して、カスタムアテンションクラス `SamVisionAttentionSplit` を作成します。`__init__` 内で、統合された `qkv` を削除し、`q`、`k`、`v` 用にそれぞれ独立した線形層を作成します。

```py
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)
```

2. `split_q_k_v_load_hook` 関数は、モデルの読み込み時に事前学習済みの `qkv` 重みを個別の `q`、`k`、`v` 重みに分割することで、あらゆる事前学習済みモデルとの互換性を確保します。

```py
    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]
```

3. `forward` パスでは、`q`、`k`、`v` がそれぞれ独立して計算されますが、アテンション機構の残りの部分は同じです。

```py
    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` のすべてのインスタンスが、分割アテンションのバージョンに置き換えられます。

`q`、`k`、`v`レイヤーがランダムに初期化されたままになるため、[from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained)を呼び出す*前に*クラスを差し替えてください。ロードフックはチェックポイントの読み込み中にのみ実行されるため、すでに読み込まれたモデルでモジュールを置き換えても、新しいレイヤーは初期化されません。

```py
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` の投影をそれぞれ独立させた場合、`q` と `v` に LoRA を適用します。

[LoraConfig](https://huggingface.co/docs/peft/package_reference/config#peft.PeftConfig) を作成し、ランク `r`、`lora_alpha`、`lora_dropout`、`task_type`、そして最も重要となるターゲットモジュールを指定します。

```py
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](https://huggingface.co/docs/peft/package_reference/config#peft.PeftConfig)を[get_peft_model](https://huggingface.co/docs/peft/package_reference/peft_model#peft.get_peft_model)に渡して、モデルにLoRAを適用します。

```py
model = get_peft_model(model, config)
```

[print_trainable_parameters](https://huggingface.co/docs/peft/package_reference/peft_model#peft.PeftMixedModel.print_trainable_parameters) を呼び出して、学習対象のパラメータ数と総パラメータ数の割合を確認します。

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