Huggingfaceの事前学習済みモデルに層を追加し，追加した層のみを学習対象とする

# 概要

以下を実現する方法の一つを紹介
- huggingfaceのモデルにAdapterやLoRA，クラス分類用のHeadといった層を追加する
    - `from_pretrained()`でロード出来た重みは`requires_grad=False`（勾配計算をしない）
    - ロード出来なかった重み（自作の層）は`requires_grad=True`（勾配計算をする）


# 実装

## 事前学習済みモデルを取得する
[Huggingface](https://huggingface.co/)では公開されているモデルを簡単に入手できます．
今回は例として`CLIPVisionModel`を用います（[ドキュメント](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPVisionModel)）．

```python
from transformers import CLIPVisionModel, CLIPVisionConfig

model_name = "openai/clip-vit-base-patch32"
config = CLIPVisionConfig.from_pretrained(model_name)
model = CLIPVisionModel.from_pretrained(
    model_name,
    config=config,
)
```

## 自作モデルを作り事前学習済みの重みをロードする
自作モデルの例として，`CLIPVisionModel`にクラス分類用の層を加えます．
なおCLIPでクラス分類を行いたい場合，Huggingfaceには`CLIPVisionModelWithProjection`が用意されているため，わざわざこのようにする必要はありませんが，説明のために作成します．
```python
import torch
import torch.nn as nn

class MyCLIPVisionModel(CLIPVisionModel):
    def __init__(self, config: CLIPVisionConfig, num_classes: int) -> None:
        super().__init__(config)
        self.head = nn.Linear(config.hidden_size, num_classes)

    def forward(
        self,
        pixel_values: torch.Tensor,
        output_attentions: bool,
        output_hidden_states: bool,
        return_dict: bool
    ) -> torch.Tensor:
        vision_outputs = super().forward(
            pixel_values,
            output_attentions,
            output_hidden_states,
            return_dict
        )
        image_embeds = vision_outputs.pooler_output
        output = self.head(image_embeds)
        return output
```
自作の`MyCLIPVisionModel`は`PreTrainedModel`を継承しているため，`from_pretrained()`を使って事前学習済みの重みを取得することが出来ます．

```python
model_name = "openai/clip-vit-base-patch32"
config = CLIPVisionConfig.from_pretrained(model_name)
num_classes = 10
model = MyCLIPVisionModel.from_pretrained(
    model_name,
    config=config,
    num_classes=num_classes
)
```


## 追加した層のみ学習対象とする

先ほどの自作クラスから`from_pretrained()`を使って事前学習済みの重みをロードした際に，以下2つの警告が出ました．（1つ目の警告は長いので省略して載せています）

- 警告1つ目
```bash
Some weights of the model checkpoint at openai/clip-vit-base-patch32 were not used when initializing MyCLIPVisionModel: ['text_model.encoder.layers.7.mlp.fc1.weight', ..., 'text_model.encoder.layers.11.self_attn.out_proj.weight']
```

CLIPというモデルはテキストエンコーダと画像エンコーダが存在します．
この警告は，「`"openai/clip-vit-base-patch32"`のcheckpointに存在する事前学習の重みを持つ層が，自作モデルには存在していない」ことを意味しています．
つまり，テキストエンコーダのcheckpointが保存されているのに使われていないという警告なので，画像エンコーダのみを使いたい場合は無視して問題ありません．

- 警告2つ目
```bash
Some weights of MyCLIPVisionModel were not initialized from the model checkpoint at openai/clip-vit-base-patch32 and are newly initialized: ['head.weight', 'head.bias']
```
ここに書かれた`['head.weight', 'head.bias']`の層は事前学習済みモデルには存在しないため，「事前学習済みモデルで初期化できない，新たに初期化した」と言ってきます（当然）．

この警告から，`['head.weight', 'head.bias']`以外の層，すなわち元の`CLIPVisionModel`の層は事前学習の重みを取得出来ていることがわかります．
しかし，このまま学習すると，全ての層が学習対象となります．


```python
# 全ての層が学習対象となっていることを確認
array = []
for param in model.parameters():
    array.append(param.requires_grad)
print(all(array)) # True
```

そこで追加した層のみ`requires_grad=True`とするため，**警告で出てくる`['head.weight', 'head.bias']`を取得します**．

[from_pretrained()のドキュメント](https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.from_pretrained)を見ると，`output_loading_info`という引数が用意されています．
これを使うと警告で出てくるパラメータ名を取得出来ます．

```python
model, loading_info = MyCLIPVisionModel.from_pretrained(
    model_name,
    config=config,
    num_classes=num_classes,
    output_loading_info=True,
)
# 警告1つ目のパラメータ名
print(loading_info['unexpected_keys']) # ['text_model.encoder.layers.7.mlp.fc1.weight', ..., 'text_model.encoder.layers.11.self_attn.out_proj.weight']
# 警告2つ目のパラメータ名
print(loading_info['missing_keys']) # ['head.weight', 'head.bias']
```
最終的には以下のようにして今回やりたかったことを実現出来ます．



```python
model_name = "openai/clip-vit-base-patch32"
config = CLIPVisionConfig.from_pretrained(model_name)
num_classes = 10
model, loading_info = MyCLIPVisionModel.from_pretrained(
    model_name,
    config=config,
    num_classes=num_classes,
    output_loading_info=True,
)
for name, param in model.named_parameters():
    if name in loading_info['missing_keys']:
        param.requires_grad = True
    else:
        param.requires_grad = False
```

# 参照

https://r0w0.hatenablog.com/entry/2020/01/29/065251

https://qiita.com/tttamaki/items/28f13a1507eb63387901