Vision Transformer：オリジナルデータでクラス分類のFineTune　と　推論モデルの作成

# Fine Tune ViT

### はじめに
HuggingFaceのTransformersと事前学習モデルが便利なので、
それらとオリジナルデータセットを抱き合わせて簡単にファインチューニングをしてみようと思います．

環境はGitにDockerFileを準備しているので、適宜使ってみてください．
> [Chamusuke/docker_vit](https://github.com/Chamusuke/docker_ViT)

一応シンプルなViT以外にも[Huggingface](https://huggingface.co/docs/transformers/model_doc/vit)には、ResNetやMobileViTなどいろいろあり、`import`を変えればFineTuneを試すことができます．ただ筆者の4GBのGPUでは、メモリ不足で全ては試せていません、、、

企業などセキュリティの厳しい方は、DockerやHuggingfaceにおいてプロキシの壁があると思います．このあたりはまた別記事でかけたらなと思います．

Colabで動かす場合、DockerFileの内容を展開していけばいいので、気が向いたら試してみます．
ローカルで動かす場合は、仮想環境に以下のライブラリをインストールすれば動くと思います。（Dockerで同じことをしているので、たぶん、、、）
```
pip install torch torchvision torchaudio --index-url　https://download.pytorch.org/whl/cu118 
pip install transformers \
    datasets \
    scikit-learn \
    evaluate \
    matplotlib \
    pipreqs \
    GPUtil \
    tensorboardX \
    requests \
    tqdm \
    regex sentencepiece sacremoses \
    && pip install -U accelerate ipywidgets \
```

※**Jetson**で試す場合は、`nvidia-smi`がサポートされていないなどGPUへのアクセスが特殊なのでベースイメージを[NVIDIA L4T PyTorch用](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/l4t-pytorch)に変更して適宜Dockerfileを調整してください．

## 全体の流れ
1. ソースコードの準備
2. データセットの準備
3. Docker：開発環境構築(jupyterlab)
4. ファインチューニング `Fine_tune_ViT.ipynb`
5. 推論の実行　        `viT_Estimete_Fine_Tune.ipynb`



## 環境
Dockerのベースイメージ：nvidia/cuda11.8 ubuntu22.04
NVIDIA GPU　T600　4GB　
python == 3.10
pytorch cuda11.8
transfomers (Hugging base)

## 準備
```
git clone https://github.com/Chamusuke/docker_ViT.git
```
ディレクトリの配置とパスについてですが、コンテナ立ち上げ後は`docker_vit`の階層が`workspace`としてマウントされるので、どこでも大丈夫です．


### データセットの準備
~~Kuggleの衛星画像写真を用いて天候の画像分類を行う~~
~~[Satellite Image Classification](https://www.kaggle.com/datasets/mahmoudreda55/satellite-image-classification)~~

分類が簡単すぎたので、鳥分類データセットでやり直しました．
[BIRDS 525 SPECIES- IMAGE CLASSIFICATION](https://www.kaggle.com/datasets/gpiosenka/100-bird-species)  

`docker_vit/data/brid_dataset/archive.zip`となるようにデータを配置し、解凍しておきます．

```shell
cd docker_vit/data 
mkdir bird_dataset
cd bird_dataset
unzip archive.zip
```

プログラム実行していくと最終的にはファイルの階層がこんな感じになります．

```:ディレクトリ階層
docker_vit/
├── Dockerfile
├── LICENSE
├── README.md
├── container_in.sh
├── data
│   ├── bird_dataset
│   │   ├── EfficientNetB0-525-(224 X 224)- 98.97.h5 #　zip解凍
│   │   ├── archive.zip
│   │   ├── birds.csv                   #　zip解凍
│   │   ├── test                        #　zip解凍
│   │   ├── train                       #　zip解凍
│   │   └── valid                       #　zip解凍
│   └── test_data
│       ├── Pembroke_Welsh_Corgi.jpg
│       ├── kiji.jpg
│       └── kingfisher.jpg
├── docker-compose.yml
├── run.sh                              #コンテナ起動用
└── src
    ├── Fine_tune_ViT.ipynb
    ├── Make_Dataset.py                 #データセット作成用（test,valid,trainに分割されていないデータを分割する）
    ├── ViT_tut.ipynb                   #推論のチュートリアル
    ├── viT_Estimete_Fine_Tune.ipynb    #FineTuneモデルをつかって推論
    └── vit-weight                      #FineTuneされたモデル
        ├── checkpoint-15500
        ├── checkpoint-16000
        ├── checkpoint-16500
        └── runs
```

#### 補足:データセットの作成
基本的なことですがファインチューニングするために、`train, valid, test`に分けておく必要があります．

※もし、用意するデータセットが以下の構造だった場合
```
dataset
├─ label_1
|   ├─pic_11.jpg
|   ├─pig_12.jpg
├─ label_2
|   ├─pic_21.jpg
|   ├─pig_22.jpg
|
```

`Make_Dataset.py`を用意したので、これで振り分けることができます．
コード内に分割する比率と参照ディレクトリ、出力ディレクトリが調整できるので適宜変更してください．

筆者は、ローカルにpythonをインストールしておらず、コンテナ内のpythonで実行したためパスを`/workspace/~`にしています．

```python: Make_Dataset.py 抜粋
# ソースとなるフォルダ
source_folder = "/workspace/data/data"
# 再構成後のフォルダ
output_folder = "/workspace/data/Dataset"

# 分割比率
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
```
`python3 Make_Dataset.py`



### Docker：開発環境の立ち上げ
#### イメージのビルド

`dcoker_vit`に移動して
```shell
docker build . -t vit_tut
```
いろいろダウンロードされると思うので気長に待ちましょう．

#### コンテナの立ち上げ & JupyterLabの起動
`bash run.sh`で実行してコンテナの立ち上げを行います．

中身はこんな感じです．
メモリとCPUの割当を行っています．適宜変更してください
-m：     20GB 
--cpus： 8個分
※dockerのバージョンによっては`runtime=NVIDIA`にする気がします．

コンテナの`/workspace`と`/docker_vit`をマウントさせています．

```shell: run.sh
docker run -it --rm --name vit_docker \
  --volume=$PWD:/workspace:rw \
  --volume=/tmp/.X11-unix:/tmp/.X11-unix:rw \
  -p 8888:8888 \
  -m 20g --cpus="8.0" \
  --gpus=all \
  vit_tut:latest \
  /bin/sh -c "jupyter-lab --no-browser --port=8888 --ip=0.0.0.0 \
    --allow-root --NotebookApp.token=''"
```

コンテナを立ち上げるとリンクが作られるのでコピペして、ブラウザでJupyterLabを開きます．
`http://localhost:8888/`



## Fine Tune & Estimate
それでは用意したデータセットをファインチューニングし、推論を実行していきます．
ここからは、` Fine_tune_ViT.ipynb`の内容と同じです．
コンテナとともに起動した**JupyterLab**で作業していきます．

### Import
今回はシンプルなVision Transformerを用います．
モデルはImageNetで1000クラス分類された事前学習済みモデルを使います．
> [google/vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224)

ここでViTのアーキテクチャを変更していろいろ試すこともできます．
用いる事前学習済みモデルとimportするImageClassificationを適宜変更してください．

**MobileViT**だとこの辺を見てコピペしてください
それだけでいろいろ試せます

> https://huggingface.co/docs/transformers/model_doc/mobilevit#transformers.MobileViTForImageClassification

```python
import requests
from PIL import Image
from tqdm import tqdm
import torch
import GPUtil
import os
from PIL import Image
from transformers import AutoImageProcessor as ViTImageProcessor

#ViTアーキテクチャに合わせる ex＞　MobileViTForImageClassification
from transformers import ViTForImageClassification as ViTForImageClassification
#from transformers import MobileViTForImageClassification as ViTForImageClassification

#　FineTuneしたモデルの保存先
weight_dir = "/workspace/src/vit-weight"
# Simple ViT model:
model_name = "google/vit-base-patch16-224" 
#Mobile_ViT model: model_name = "apple/mobilevit-small"

device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
torch.cuda.empty_cache()
GPUtil.showUtilization()
```

### 事前学習済みモデルのダウンロード
プロキシがある場合はproxy_dic={http: < your proxy >}の辞書形式で渡してください．コンテナのプロキシが通っている場合もしかしたら必要ないかもしれません．
`ViTImageProcessor.from_pretrained(model_name,proxy_dic=proxy_dic)`

```python
image_processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)
```

### データセットの読み込みと確認
#### データセットの読み込み
用意したデータセットを読み込みます．
```python
from datasets import load_dataset
ds = load_dataset("imagefolder", data_dir="/workspace/data/bird_dataset/")
ds
```

`train`,`validation`,`test`に振り分けられていたらOKです．
```:output
DatasetDict({
    train: Dataset({
        features: ['image', 'label'],
        num_rows: 84635
    })
    validation: Dataset({
        features: ['image', 'label'],
        num_rows: 2625
    })
    test: Dataset({
        features: ['image', 'label'],
        num_rows: 2625
    })
})
```

#### ラベルの確認
データセットのtrain内のクラスが1つの場合、Noneと表示されます． 
このときラベルが必要であるなら、データセットの読み込み時に`drop_labels=False`のオプションを追加してラベルをつけることができます．
`load_dataset("imagefolder", data_dir="/workspace/data/Dataset", drop_labels=False)`

大量のラベルが出てくるので、今回は102番目と202番目の画像に対応するラベルを確認してみます．
```python
labels = ds["train"].features["label"]
#print(labels)
print(labels.int2str(ds["train"][102]["label"]))
print(labels.int2str(ds["train"][202]["label"]))
```
```:output
ABBOTTS BABBLER
ABBOTTS BOOBY
```

#### データセットとラベルの関係確認
読み込んでいる画像データとラベルの確認をします．
```python
import random
import matplotlib.pyplot as plt

def show_image_grid(dataset, split, grid_size=(4,4)):
    indices = random.sample(range(len(dataset[split])), grid_size[0]*grid_size[1])
    images = [dataset[split][i]["image"] for i in indices]
    labels = [dataset[split][i]["label"] for i in indices]
    
    fig, axes = plt.subplots(nrows=grid_size[0], ncols=grid_size[1], figsize=(8,8))
    for i, ax in enumerate(axes.flat):
        ax.imshow(images[i])
        ax.axis('off')
        ax.set_title(ds["train"].features["label"].int2str(labels[i]))
    
    plt.show()

show_image_grid(ds, "test")
```
![Image_output](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3008380/b9d4599d-1a49-cdaf-6ca6-aad51e88997f.png)

### 前処理
画像の読み込みがうまくできたことを確認できたら、すべてのRGB画像を244*244（使うモデルによって変化します）のPytorchテンソルに変換します．
```python
def transform(examples):
  inputs = image_processor([img.convert("RGB") for img in examples["image"]], return_tensors="pt")
  inputs["labels"] = examples["label"]
  return inputs

dataset = ds.with_transform(transform)
```

#### テンソルの確認
```python
for item in dataset["train"]:
  print(item["pixel_values"].shape)
  break
```
```:output
torch.Size([3, 224,224])　ー＞　RGB 3チャンネル　
```

#### バッチを積み重ねて照合させるための関数
データセットからのバッチを受け取り、ピクセル値とラベルを含む辞書形式のテンソルを返す関数を作ります
```python
def collate_fn(batch):
  return {
      "pixel_values": torch.stack([x["pixel_values"] for x in batch]),
      "labels": torch.tensor([x["labels"] for x in batch]),
  }
```

### モデル学習の準備
ファインチューニングするためのパラメータ等の準備を行います．

#### メトリックを定義
TrainerAPIが、`compute_metric`関数を用いるため用意します．
今回はAccuracyとF1の算出を行っています．
evaluateを用いましたが、scikitlearnを用いても良いです．
場合によっては、自分でごりごり作るほうが良いのかもしれません
evaluateのモジュールは結構エラーが出ます:pray:

（eval_pred.predicions, axis=1）で各クラスにおける最大確率を算出し、argmax()でそのindexを求めています．  
そこからラベルを取り出してます．
```python
from evaluate import load
import numpy as np

accuracy = load("accuracy")
f1 = load("f1")

def compute_metrics(eval_pred):
  accuracy_score = accuracy.compute(predictions=np.argmax(eval_pred.predictions, axis=1), references=eval_pred.label_ids)
  f1_score = f1.compute(predictions=np.argmax(eval_pred.predictions, axis=1), references=eval_pred.label_ids, average="macro")
  return {**accuracy_score, **f1_score}
  ```

#### 出力クラスの初期化
 事前学習済みモデルはImageNetで1000クラスの出力が設定されています．
そこで、オリジナルデータセット用に出力し直すため出力ヘッドを初期化します．厳密には少し違いますが、これがFineTuneと言うやつです．
筆者はこの設定をよく忘れてしまい、学習したはずなのに精度が全く上がらない現象を何度も味わいました、、、
  
```python
model = ViTForImageClassification.from_pretrained(
    model_name,
    num_labels=len(labels),
    id2label={str(i): c for i, c in enumerate(labels)},
    label2id={c: str(i) for i, c in enumerate(labels)},
    ignore_mismatched_sizes=True,
)
```

#### 学習パラメータの設定
ここからは、よく見るパラメータたちです．
筆者の４GBGPUだとバッチサイズは以下のようにショボいです．
皆さんのスーパーPCに合わせて変更してください．

```python
from transformers import TrainingArguments

training_args = TrainingArguments(
  output_dir=weight_dir, # Fine tuneモデルの出力先
  per_device_train_batch_size=10, # バッチサイズ
  per_device_eval_batch_size=8,   # バッチサイズ
  evaluation_strategy="steps",    # ステップごとに評価
  num_train_epochs=2,             # epoch数
  #fp16=True,                      # 混合精度
  save_steps=500,                  # チェックポイントを保存頻度
  eval_steps=500,                  # モデルの評価を行う頻度
  logging_steps=500,               # 更新ステップごとにログ
  # save_steps=50,
  # eval_steps=50,
  # logging_steps=50,
  save_total_limit=3,             # ディスク上のチェックポイントの総量を制限
  remove_unused_columns=False,    # データセットから未使用の列を削除する
  push_to_hub=False,              # モデルをハブにプッシュしない
  report_to='tensorboard',        # メトリクスをTensorBoardに渡す
  load_best_model_at_end=True,    #6_full_eval (bool, optional, defaults to False) — Whether to use full float16 evaluation instead of 32-bit. This will be faster and save memory but can harm metric values. 学習の最後に最良のモデルをロードする
  use_cpu=False,                   # CPUの使用　Flase->GPU・TPUの使用
  #fp16_full_eval=True
)
```

#### 学習用のクラス
設定したパラメータやデータセット、評価関数を学習用のクラスに渡します．
```python
from transformers import Trainer

trainer = Trainer(
    model=model,                        # 事前学習済みモデルの用意
    args=training_args,                 # 設定したArgumentsを渡す
    data_collator=collate_fn,           # ラベル照合用の関数を渡す
    compute_metrics=compute_metrics,    # 評価時に使うメトリック関数を渡す
    train_dataset=dataset["train"],     # 訓練データセット
    eval_dataset=dataset["validation"], # 評価データセット
    tokenizer=image_processor,          # トークナイザー
)
```

### ファインチューニング
長くなりましたようやくモデルの学習準備が整いました．
以下のコードを実行して気長に待ちましょう．
ecpoch=2も必要ないのではという伸びです．
ViTすごい、、、
GPU、メモリも頑張ってくれました．
```python
trainer.train()
GPUtil.showUtilization()
```
```:output
 [16928/16928 5:05:30, Epoch 2/2]
Step	Train Loss	Valid Loss	Accuracy	F1
500	    5.134500	3.646356	0.645333	0.591693
1000	2.649400	1.659642	0.862095	0.838723
1500	1.249600	0.773362	0.930667	0.921018
.....   ........    ........    ........    ........
16000	0.059600	0.083765	0.986286	0.984195
16500	0.057400	0.084243	0.985905	0.983811

| ID | GPU | MEM |
------------------
|  0 | 98% | 87% |
```

## FineTuneモデルの検証
### テストデータで確認
ここまで残しておいた`test`を用いてモデルの検証します．
```python
output= trainer.predict(dataset["test"])
output.metrics
```
```:output
{'test_loss': 0.02646920271217823,
 'test_accuracy': 0.993904761904762,
 'test_f1': 0.993770081770082,
 'test_runtime': 97.5077,
 'test_samples_per_second': 26.921,
 'test_steps_per_second': 3.374}
 ```

 ### 混合行列
 混合行列のヒートマップで確認してみます．
 ただ今回のデータセットだと分類クラスが多すぎてわかりにくいものとなりました．少ないクラスだと見やすいかなと思います

横軸が推論ラベルで縦が正解ラベルです．対角行列担っていればいいモデルなのですが、、、見た感じだとなんか良さそうです．
 
 ```python
 import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix

y_preds = np.argmax(output.predictions, axis=1)
y_valid = np.array(ds["test"]["label"])
labels = ds["train"].features["label"].names

def plot_confusion_matrix(y_preds, y_true, labels):
    cm = confusion_matrix(y_true, y_preds, normalize="true")
    fig, ax = plt.subplots(figsize=(16,16))
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
    disp.plot(cmap="Blues", values_format=".2f", ax=ax, colorbar=False)
    ax.set_xticklabels(labels, rotation=90) 
    ax.set_yticklabels(labels, rotation=0) 
    plt.title("Normalized confusion matrix")
    plt.show()

plot_confusion_matrix(y_preds, y_valid, labels[:-1])
```
![混合行列](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3008380/1455a200-981a-7e94-1475-d0be7c41993b.png)


## FineTuneモデルで推論
モデルの作成ができたので、作成したオリジナルモデルで推論を実行します．
移植性も考慮して、新しいファイルを用意してください．
ここからは`viT_Estimete_Fine_Tune.ipynb`になります．

補足：GPUのメモリエラーを回避するため、`Fine_tune_ViT.ipynb`のカーネルrestartしておくと良いです

### 画像の準備
推論したい画像を適宜用意します．
```python
import os
from PIL import Image
import requests
import matplotlib.pyplot as plt
import torch
import GPUtil


print("cuda" if torch.cuda.is_available() else "cpu")
torch.cuda.empty_cache()
GPUtil.showUtilization()

#file_path = '/workspace/data/test_data/cher_1.jpg'
file_path = '/workspace/data/bird_dataset/test/PYGMY KINGFISHER/3.jpg'
image = Image.open(file_path)
image
```
```:output
cuda
| ID | GPU | MEM |
------------------
|  0 |  0% | 40% |
```
![test_Image](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3008380/353b83f3-e81b-0a51-ea58-552974cddd45.png)

### モデルの準備
ファインチューニングして保存されたオリジナルモデルを参照させてください．

```python
import torch
from transformers import ViTForImageClassification, AutoImageProcessor
import time

# デバイスの設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

#モデルのロード
best_checkpoint = 16500
model_id = f"/workspace/src/vit-weight/checkpoint-{best_checkpoint}"
model_gpu = ViTForImageClassification.from_pretrained(model_id).to(device)
image_processor_gpu =  AutoImageProcessor.from_pretrained(model_id)
```




### 推論の実行①
推論を実行し、クラス分類を実行します．
```python
len = 100 
start = time.time()

for _ in range(len):
    inputs_gpu = image_processor_gpu(images=image, return_tensors='pt').to(device)
    
    with torch.no_grad():
        outputs_gpu = model_gpu(**inputs_gpu)
        logits_gpu = outputs_gpu.logits 
        predicted_idx_gpu = logits_gpu.argmax(-1).item() 

print('class:', model_gpu.config.id2label[predicted_idx_gpu]) 
fps = len/(time.time() - start)
print("FPS:", fps)
GPUtil.showUtilization()
```
```:output
class: PYGMY KINGFISHER
FPS: 21.76284647777463
| ID | GPU | MEM |
------------------
|  0 | 92% | 32% |
```


### 推論の実行②
クラス分類において、確率が最大となるクラスが返されます．ここでは、大きいものから3つ取り出してみます．

いい具合に識別できています．
```python
def get_prediction_probs(model, url_or_path, num_classes=3):
    img = Image.open(url_or_path)
    pixel_values = image_processor_gpu(img, return_tensors="pt")["pixel_values"].to(device)
    output = model_gpu(pixel_values)
    probs, indices = torch.topk(output.logits.softmax(dim=1), k=num_classes)
    id2label = model_gpu.config.id2label
    classes = [id2label[idx.item()] for idx in indices[0]]
    probs = probs.squeeze().tolist()
    results = dict(zip(classes, probs))
    return results
```
```python
get_prediction_probs(model_id, file_path)
```
```:output
{'PYGMY KINGFISHER': 0.9989274144172668,
 'MALACHITE KINGFISHER': 0.0003117609303444624,
 'RUFOUS KINGFISHER': 0.00014921906404197216}
```

## 参考
[Vision Transformer (ViT)](https://huggingface.co/docs/transformers/model_doc/vit)
[How to Fine Tune ViT for Image Classification using Transformers in Python](https://thepythoncode.com/article/finetune-vit-for-image-classification-using-transformers-in-python)