Fine Tune ViT
はじめに
HuggingFaceのTransformersと事前学習モデルが便利なので、
それらとオリジナルデータセットを抱き合わせて簡単にファインチューニングをしてみようと思います．
環境はGitにDockerFileを準備しているので、適宜使ってみてください．
一応シンプルな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を調整してください．
全体の流れ
- ソースコードの準備
- データセットの準備
- Docker：開発環境構築(jupyterlab)
- ファインチューニング
Fine_tune_ViT.ipynb - 推論の実行
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/~にしています．
# ソースとなるフォルダ
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をマウントさせています．
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クラス分類された事前学習済みモデルを使います．
ここでViTのアーキテクチャを変更していろいろ試すこともできます．
用いる事前学習済みモデルとimportするImageClassificationを適宜変更してください．
MobileViTだとこの辺を見てコピペしてください
それだけでいろいろ試せます
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です．
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"]))
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
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()
[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
{'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
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()
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)
{'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