Transformers documentation
モデルコンポーネントのカスタマイズ
モデルコンポーネントのカスタマイズ
モデルをカスタマイズするもう一つの方法は、新しいモデルをゼロから書くのではなく、既存のコンポーネントを変更することです。これにより、特定のユースケースに合わせてモデルを調整できます。例えば、新しいレイヤーを追加したり、アーキテクチャの注意機構を最適化したりできます。カスタマイズは直接 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")
Attentionクラス
Segment Anythingは画像セグメンテーションモデルであり、そのアテンション機構においてクエリ・キー・バリュー（qkv）の射影を結合しています。学習可能なパラメータ数と計算オーバーヘッドを削減するために、qkv射影にLoRAを適用できます。これには、qkv射影を分割して、qとvを個別にLoRAのターゲットとして指定できるようにする必要があります。
SamVisionAttentionクラスをサブクラス化して、カスタムアテンションクラスSamVisionAttentionSplitを作成します。__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 のすべてのインスタンスが、分割アテンションのバージョンに置き換えられます。
q、k、vレイヤーがランダムに初期化されたままになるため、from_pretrained()を呼び出す前にクラスを差し替えてください。ロードフックはチェックポイントの読み込み中にのみ実行されるため、すでに読み込まれたモデルでモジュールを置き換えても、新しいレイヤーは初期化されません。
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 を作成し、ランク 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"