本記事の要約

　前回ビジョン系モデルのローカル環境構築について書かせてもらいましたが、特にTransformersを使った推論とファインチューニングのコードについては詳細をご紹介できませんでした。そこで今回はMicrosoft社の優れたVLMであるFlorence-2のPythonコードを紹介させて頂き、何かの参考にして頂ければと思います。

　最初にモデルを選んだ背景や狙いについて説明しておきます。
１. なんでFlorence-2なのか？
　VLMには特に最近色々優秀なモデルが出てきてますが、日本語で色々なタスクを処理できるモデルはそう多くなく、また利用に際してのポリシーやライセンスや利用規約の縛りも様々なので、自由にチューニングしてローカルで安心して使えるモデルはまだ少ないと考えてます。そんな中で、前回記事でご紹介した幾つかのモデルは全て条件を満たすものはないものの、近い存在ではないかと思われます。
　Florence-2のメリットは、サイズが小さく(base-ftモデルのfp16でSafetensorsが542MB)GPUやOllamaが無くても十分ローカルで使えること、9種類の画像認識タスクがゼロショットで使えてその性能が高く、更にライセンスがMITライセンスなので扱いやすい事です。残念なのは日本語未対応な点ですが、今回ご紹介するように独自で日本語化のファインチューニングも可能なので、これで何とか条件をクリア出来るのではと考えました。
２. 他のモデルはどうなのか？
前回ご紹介したllava-jp-1.3b-v1.1は日本語能力も高く魅力的なモデルですが、サイズの関係でGGUF化＋量子化したかったのですが、出来なかったので諦めました。GPUとメモリに余裕がある方はトライしてみて下さい。nanoLLaVA-1.5やmoondream2はOllamaで動かせるメリットはありますが、日本語化のファインチューニングする度にGGUF化＋量子化するのは辛いと考え保留としました。

1. Florence-2の概要

　Microsoft社（以下MS社）がMITライセンスで公開したVLM（MS社によるとVFMらしい：ビジョン基盤モデル）で、ピクセル・リージョン単位のタスク向けで、現実の物体を含むシナリオ向けのモデルです。ゼロショットで下記のような９種類のタスクに対応してます（HuggingFace公式での説明。更に数種類のタスクが可能という説明も存在する）。
・CAPTION：キャプション（画像の説明）
・DETAILED_CAPTION：詳細なキャプション
・MORE_DETAILED_CAPTION：更により詳細なキャプション
・CAPTION_TO_PHRASE_GROUNDING：入力されたテキストを関連する画像領域に紐づける
・OD：オブジェクトディテクション（物体検出）、検出された領域の座標を出力
・DENSE_REGION_CAPTION：高密度領域キャプション
・REGION_PROPOSAL：領域提案
・OCR：OCR（テキスト抽出）
・OCR_WITH_REGION：領域検出OCR

　今回は日本語が必須と考えられるCAPTIONとOCRの日本語化ファインチューニングにトライしました。

2. Florence-2のファインチューニング

　以下の参考記事を元に色々試作錯誤した結果、以下のコードでColabやローカルMacの環境で共通して動作するファインチューニングに成功しました。
まずPytorch等はすでにインストール済みの前提で、必要なパッケージをインストします。

pip install -q transformers datasets

以下が実際のファインチューニングのコードです。

# モデルの読み込み
from transformers import AutoProcessor, AutoModelForCausalLM
import torch

CHECKPOINT = "microsoft/Florence-2-base-ft"
REVISION = 'refs/pr/6'

#Apple siliconの場合
#device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
#その他
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

model = AutoModelForCausalLM.from_pretrained(CHECKPOINT, trust_remote_code=True, revision=REVISION).to(device)
processor = AutoProcessor.from_pretrained(CHECKPOINT, trust_remote_code=True, revision=REVISION)

# datasetの読み込み
from datasets import load_dataset

# CAPTIONの場合
dataset_name = "Kendamarron/japanese-photo-instruction"
# OCR_JPの場合は上記をコメントアウト＆下記のコメント削除
#dataset_name = "EtashGuha/JapaneseDocQA"

dataset = load_dataset(dataset_name, split="train")
### Filterを挿入する場所　###

print(len(dataset))
dataset = dataset.train_test_split(train_size=0.9, test_size=0.1, shuffle=True, seed=42)

""" <OCR_JP>の場合は以下を追加
# <OCR_JP>を新たな専用タスクとしてカスタマイズ
import torch.nn as nn
model.ocr_jp = nn.Sequential(
    nn.Linear(768, 512),  # Adjusting the input features for the new task
    nn.ReLU(),           # Adding non-linearity
    nn.Linear(512, 10)  # Output layer for the number of classes
)
"""

# OCRDatasetクラスを定義
import torch
torch.cuda.empty_cache()

from torch.utils.data import Dataset

class OCRDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        #print(len(self.data))
        try:
            example = self.data[idx]
            question = "<CAPTION>"  # <OCR_JP>は"<OCR_JP>"に置き換える
            answer = example["caption"] # <OCR_JP>はexample['text']に置き換える
            #print(question,answer)
            image = example['image']
            if image.mode !="RGB":
                image = image.convert("RGB")
            return question, answer, image
        except KeyError as e:
            print(f"KeyError encountered at index {idx}: {e}")
            return "", "", None

# データセットの作成
train_dataset = OCRDataset(dataset['train'])
eval_dataset = OCRDataset(dataset['test'])

# collate関数の定義
IGNORE_ID = -100 # Pytorch ignore index when computing loss
MAX_LENGTH = 1024
def collate_fn(batch):
    questions, answers, images = zip(*batch)
    inputs = processor(text=list(questions), images=list(images), return_tensors="pt", padding=True, ).to(device)
    input_ids = inputs["input_ids"]
    #pixel_values = inputs["pixel_values"]
    labels = processor.tokenizer(text=list(answers), return_tensors="pt", padding=True, return_token_type_ids=False, truncation=True, max_length=1024).input_ids.to(device)
    labels[labels == processor.tokenizer.pad_token_id] = IGNORE_ID
    #return_data = {"input_ids": input_ids, "labels": labels}
    return_data = {**inputs, "labels": labels}
    return return_data
# loraの場合に上記collate_fnでエラーになる場合は下記で試して下さい
"""
def collate_fn(batch):
    questions, answers, images = zip(*batch)
    inputs = processor(text=list(questions), images=list(images), return_tensors="pt", padding=True).to(device)
    return inputs, answers
"""

#TrainingArgumentsの設定
from transformers import TrainingArguments

args=TrainingArguments(
    output_dir="./Florence-2-trainer-OCRJP",
    num_train_epochs=10, # 自分の環境に合わせて設定して下さい
    learning_rate=1e-4, # 自分の環境に合わせて設定して下さい
    per_device_train_batch_size=2, # 自分の環境に合わせて設定して下さい
    per_device_eval_batch_size=2, # 自分の環境に合わせて設定して下さい
    gradient_accumulation_steps=4, # 自分の環境に合わせて設定して下さい
    save_strategy="epoch",
    logging_strategy="steps",
    logging_steps = 200, # 自分の環境に合わせて設定して下さい
    eval_strategy="steps",
    eval_steps = 200, # 自分の環境に合わせて設定して下さい
    save_total_limit=5, # 自分の環境に合わせて設定して下さい
    load_best_model_at_end=False,
    label_names=["labels"],
    report_to="tensorboard",
    remove_unused_columns=False,  # needed for data collator
    dataloader_pin_memory=False,
)

# trainerの定義
from transformers import Trainer

trainer = Trainer(
    model=model,
    tokenizer=processor,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=collate_fn, # dont forget to add custom data collator
    args=args
)

# vision towerパラメータを凍結し処理を軽減
for param in model.vision_tower.parameters():
  param.requires_grad = False

model_total_params = sum(p.numel() for p in model.parameters())
model_train_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"Number of trainable parameters {model_train_params} out of {model_total_params}, rate: {model_train_params/model_total_params:0.3f}")

# ファインチューニングの実行
trainer.train()

# HugginFaceへ結果をアップロード
trainer.push_to_hub("<your hf_repo_id>", token="<your hf_token>") # <your hf_repo_id>と<your hf_token>を設定して下さい

上記Pythonコードを実行すると下記のような学習結果が表示されます。

Step	Training Loss	Validation Loss
200	1.860700	1.257793
400	1.013400	0.993707
600	0.689100	0.873091
800	0.492400	0.816602
1000	0.347600	0.792435

　基本的な処理の流れは、前回記事に記載した通りですが、ベースはCAPTION用のコードでOCR_JPに変更する場合は適宜必要箇所を入れ替えて使って下さい。また、datasetのサイズを調整したい場合はデータを読み込んだ後「### Filterを挿入する場所　###」に下記のようなフィルタを入れて調整できます。

dataset = dataset.filter(lambda example: len(str(example["caption"])) < 500 and len(str(example["caption"])) > 100)

3. Florence-2のloraによるファインチューニング

　loraによるファインチューニングは上記コードに下記のようにpeftによるlora設定を追加します。上記コードの#TrainingArgumentsの設定以下を下記のコードに置き換えて下さい。

from peft import LoraConfig, get_peft_model

peft_config = LoraConfig(
    r=128, # 自分の環境に合わせて設定して下さい
    lora_alpha=256, # 自分の環境に合わせて設定して下さい
    target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "linear", "Conv2d", "lm_head", "fc2"],
    # <OCR_JP>は、modules_to_save=["ocr_jp"]を追加する
    task_type="CAUSAL_LM",
    lora_dropout=0.05,
    bias="none",
    inference_mode=False,
    use_rslora=True,
    init_lora_weights="gaussian",
    revision=REVISION
)

lora_model = get_peft_model(model, peft_config)
lora_model.to(device)

# vision towerパラメータを凍結し処理を軽減
for param in lora_model.vision_tower.parameters():
  #param.is_trainable = False
  param.requires_grad = False
lora_model.print_trainable_parameters()

# train関数の定義
def train_model(train_loader, val_loader, model, processor, epochs=10, lr=1e-6):
    optimizer = AdamW(model.parameters(), lr=lr)
    num_training_steps = epochs * len(train_loader)
    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=0,
        num_training_steps=num_training_steps,
    )

    for epoch in range(epochs):
        model.train()
        train_loss = 0
        for inputs, answers in tqdm(train_loader, desc=f"Training Epoch {epoch + 1}/{epochs}"):
            input_ids = inputs["input_ids"]
            pixel_values = inputs["pixel_values"]
            labels = processor.tokenizer(
                text=answers,
                return_tensors="pt",
                padding=True,
                return_token_type_ids=False
            ).input_ids.to(device)
            outputs = model(input_ids=input_ids, pixel_values=pixel_values, labels=labels)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            train_loss += loss.item()
        avg_train_loss = train_loss / len(train_loader)
        print(f"Average Training Loss: {avg_train_loss}")
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for inputs, answers in tqdm(val_loader, desc=f"Validation Epoch {epoch + 1}/{epochs}"):
                input_ids = inputs["input_ids"]
                pixel_values = inputs["pixel_values"]
                labels = processor.tokenizer(
                    text=answers,
                    return_tensors="pt",
                    padding=True,
                    return_token_type_ids=False
                ).input_ids.to(device)
                outputs = model(input_ids=input_ids, pixel_values=pixel_values, labels=labels)
                loss = outputs.loss
                val_loss += loss.item()
            avg_val_loss = val_loss / len(val_loader)
            print(f"Average Validation Loss: {avg_val_loss}")
        output_dir = save_dir+"/epoch_{epoch+1}"
        model.save_pretrained(output_dir, save_embedding_layers=True)
        processor.save_pretrained(output_dir)

# train_modelの実行
EPOCHS = 10 # 自分の環境に合わせて設定して下さい
LR = 7e-5 # 自分の環境に合わせて設定して下さい
train_model(train_loader, val_loader, lora_model, processor, epochs=EPOCHS, lr=LR)

# loraモデルの保存
lora_model.save_pretrained(save_dir, save_embedding_layers=True)
processor.save_pretrained(save_dir)

# HuggingFaceへBFLOAT16でアップロードする例
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(CHECKPOINT, low_cpu_mem_usage=True, return_dict=True, trust_remote_code=True, torch_dtype=torch.bfloat16)
peft_model = PeftModel.from_pretrained(base_model, save_dir).to(device)
merged_model = peft_model.merge_and_unload(safe_merge=True) #add safe_merge=True
merged_model.push_to_hub("<your hf_repo_id>", token="<your hf_token>", variant="bf16")
processor.push_to_hub("<your hf_repo_id>", token="<your hf_token>")

　以上です。お疲れ様でした。
なお、CAPTIONはFlorence2のデフォルトの<CAPTION>でInference出来ますが、OCR_JPは独自に追加した<OCR_JP>をpromptに設定して使って下さい。

ファインチューニングしたサンプルモデル：
・CAPTION：　aipib/Florence-2-VQAJP2
・OCR_JP:　aipib/Florence-2-OCRJP-remake4

参考記事：