PEFTライブラリを使用して自由にLoRAチューニングする方法

2023年6月17日
機械学習

目次

ベースモデルのパラメータを確認する

以下のようなコードでベースモデルのパラメータを表示して確認する。


from transformers import AutoModel

model = AutoModel.from_pretrained("cl-tohoku/bert-base-japanese-v3")
for name, param in model.named_parameters():
    print(name)

表示されたパラメータのうち、LoRAを適用したいパラメータの名前を把握しておく。
今回は以下のような層に適用するものとして考える。

encoder.layer.X.attention.self.query.weight
encoder.layer.X.attention.self.query.bias
encoder.layer.X.attention.self.key.weight
encoder.layer.X.attention.self.key.bias
encoder.layer.X.attention.self.value.weight
encoder.layer.X.attention.self.value.bias

モデルを作成する

続いて、ベースモデルを利用して自身のやりたいタスク用のモデルを構築する。
ここでは、特に特別なことは必要ない。
ただし、自分が追加した層のうち学習させる必要がある層の名前は覚えておく。
ここでは、my_classifier層を学習させる。

class BertClassifier(nn.Module):
    def __init__(self):
        super(BertClassifier, self).__init__()
        self.config = BertConfig.from_pretrained('cl-tohoku/bert-base-japanese-v3')
        self.bert = BertModel.from_pretrained('cl-tohoku/bert-base-japanese-v3',
                                              output_attentions=True,
                                              output_hidden_states=True)

        self.my_classifier = nn.Linear(768*1, 3)

        # 重み初期化処理
        nn.init.normal_(self.my_classifier.weight, std=0.02)
        nn.init.normal_(self.my_classifier.bias, 0)

    # clsトークンのベクトルを取得する用の関数を用意
    def _get_cls_vec(self, vec):
        return vec[:,0,:].view(-1, 768)

    def forward(self, input_ids):

        # 順伝播の出力結果は辞書形式なので、必要な値のkeyを指定して取得する
        output = self.bert(input_ids)
        attentions = output['attentions']
        hidden_states = output['hidden_states']

        vec = self._get_cls_vec(hidden_states[-1])

        # 全結合層でクラス分類用に次元を変換
        out = self.my_classifier(vec)
    
        return F.log_softmax(out, dim=1), attentions

PEFTライブラリとかのバージョンによってはうまく動かないので以下のようにしてみるのもあり…

class BertClassifier(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)

        self.bert = BertModel(config)
        self.my_classifier = nn.Linear(768*1, 3)

        # 重み初期化処理
        nn.init.normal_(self.my_classifier.weight, std=0.02)
        nn.init.normal_(self.my_classifier.bias, 0)

    # clsトークンのベクトルを取得する用の関数を用意
    def _get_cls_vec(self, vec):
        return vec[:,0,:].view(-1, 768)

    def forward(self, input_ids):

        # 順伝播の出力結果は辞書形式なので、必要な値のkeyを指定して取得する
        output = self.bert(input_ids)
        attentions = output['attentions']
        hidden_states = output['hidden_states']

        vec = self._get_cls_vec(hidden_states[-1])

        # 全結合層でクラス分類用に次元を変換
        out = self.my_classifier(vec)
    
        return F.log_softmax(out, dim=1), attentions

以下のように読み込める

model = BertClassifier.from_pretrained('cl-tohoku/bert-base-japanese-v3')

LoRAの設定を行う

肝となるのはtarget_modulesとmodules_to_saveの部分である。
target_modulesにはLoRAチューニングをしたい層の名前を入れる。
modules_to_saveにはLoRA層以外の学習させたい層の名前を入れる。
もちろん、target_modulesとmodules_to_saveに入れる層はベースモデルの層・自分が追加した層をどのように組み合わせても良い。
層の指定には正規表現とか使えるらしい…

model = BertClassifier()

config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=['query','key','value'],
    lora_dropout=0.05,
    bias="none",
    modules_to_save=['my_classifier']
)
model = get_peft_model(model, config)

パラメータの保存

model.save_pretrained("model_dir_path")

パラメータの読み込み

model = BertClassifier()
model = PeftModel.from_pretrained(model, "model_dir_path")