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

Fine Tune ViT

はじめに

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

環境はGitにDockerFileを準備しているので、適宜使ってみてください．

Chamusuke/docker_vit

一応シンプルなViT以外にもHuggingfaceには、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用に変更して適宜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

分類が簡単すぎたので、鳥分類データセットでやり直しました．
BIRDS 525 SPECIES- IMAGE CLASSIFICATION

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

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/~にしています．

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に移動して

docker build . -t vit_tut

いろいろダウンロードされると思うので気長に待ちましょう．

コンテナの立ち上げ & JupyterLabの起動

bash run.shで実行してコンテナの立ち上げを行います．

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

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

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

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

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

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

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)

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

データセットの読み込みと確認

データセットの読み込み

用意したデータセットを読み込みます．

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番目の画像に対応するラベルを確認してみます．

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

データセットとラベルの関係確認

読み込んでいる画像データとラベルの確認をします．

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")

前処理

画像の読み込みがうまくできたことを確認できたら、すべてのRGB画像を244*244（使うモデルによって変化します）のPytorchテンソルに変換します．

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)

テンソルの確認

for item in dataset["train"]:
  print(item["pixel_values"].shape)
  break

output

torch.Size([3, 224,224])　ー＞　RGB 3チャンネル　

バッチを積み重ねて照合させるための関数

データセットからのバッチを受け取り、ピクセル値とラベルを含む辞書形式のテンソルを返す関数を作ります

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のモジュールは結構エラーが出ます

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

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と言うやつです．
筆者はこの設定をよく忘れてしまい、学習したはずなのに精度が全く上がらない現象を何度も味わいました、、、

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に合わせて変更してください．

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
)

学習用のクラス

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

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、メモリも頑張ってくれました．

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を用いてモデルの検証します．

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}

混合行列

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

横軸が推論ラベルで縦が正解ラベルです．対角行列担っていればいいモデルなのですが、、、見た感じだとなんか良さそうです．

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])

FineTuneモデルで推論

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

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

画像の準備

推論したい画像を適宜用意します．

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% |

モデルの準備

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

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)

推論の実行①

推論を実行し、クラス分類を実行します．

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つ取り出してみます．

いい具合に識別できています．

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

get_prediction_probs(model_id, file_path)

output

{'PYGMY KINGFISHER': 0.9989274144172668,
 'MALACHITE KINGFISHER': 0.0003117609303444624,
 'RUFOUS KINGFISHER': 0.00014921906404197216}

参考

Vision Transformer (ViT)
How to Fine Tune ViT for Image Classification using Transformers in Python