Transformers で Swin Transformer v2 をファインチューンして画像分類

from datasets import load_dataset, DatasetDict
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    DefaultDataCollator,
    TrainingArguments,
    Trainer,
)
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
import evaluate
import numpy as np
import wandb


CACHE_DIR = "/huggingface/cache" # これは個人的に使ってるだけ
MODEL_NAME = "microsoft/swinv2-base-patch4-window12-192-22k"
OUTPUT_NAME = 出力モデル名(と、wandbのプロジェクト名)
HUB_REPO_NAME = HuggingFaceのレポ名(アップするなら)


processor = AutoImageProcessor.from_pretrained(
    MODEL_NAME,
    cache_dir=CACHE_DIR,
)

normalize = Normalize(mean=processor.image_mean, std=processor.image_std)

size = (
    processor.size["shortest_edge"]
    if "shortest_edge" in processor.size
    else (processor.size["height"], processor.size["width"])
)

_transforms = Compose([RandomResizedCrop(size), ToTensor(), normalize])


def transforms(examples):
    examples["pixel_values"] = [
        _transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples


dataset = load_dataset("imagefolder", data_dir="./dataset")
assert isinstance(dataset, DatasetDict)
dataset["train"] = dataset["train"].with_transform(transforms)
dataset["test"] = dataset["test"].with_transform(transforms)
print(dataset)
labels = dataset["train"].unique("label")


data_collator = DefaultDataCollator()

accuracy = evaluate.load("accuracy")


def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)


id2label = {
    0: "best quality",
    1: "high quality",
    2: "medium quality",
    3: "low quality",
    4: "worst quality",
    5: "parts",
    6: "other",
}
label2id = {v: k for k, v in id2label.items()}

model = AutoModelForImageClassification.from_pretrained(
    MODEL_NAME,
    ignore_mismatched_sizes=True,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
    cache_dir=CACHE_DIR,
)

wandb.init(project=OUTPUT_NAME)
training_args = TrainingArguments(
    output_dir=OUTPUT_NAME,
    remove_unused_columns=False,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    optim="adafactor",
    lr_scheduler_type="cosine",
    learning_rate=5e-5,
    auto_find_batch_size=True,
    gradient_accumulation_steps=4,
    num_train_epochs=10,
    save_total_limit=2,
    save_safetensors=True,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=True,
    hub_model_id=HUB_REPO_NAME,
    hub_private_repo=True,
    report_to=["wandb"],
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=processor,
    compute_metrics=compute_metrics,
)

trainer.train()

trainer.evaluate()

trainer.save_model(OUTPUT_NAME)
trainer.push_to_hub(HUB_REPO_NAME)

python train_swinv2.py

from transformers import AutoImageProcessor
import torch
from PIL import Image

from transformers import AutoModelForImageClassification

import gradio as gr

CACHE_DIR = "/huggingface/cache"
MODEL_NAME = 学習したモデルのパス or  レポ名

image_processor = AutoImageProcessor.from_pretrained(
    MODEL_NAME,
    cache_dir=CACHE_DIR,
)
model = AutoModelForImageClassification.from_pretrained(
    MODEL_NAME,
    cache_dir=CACHE_DIR,
)


def classify_image(image: Image.Image):
    inputs = image_processor(image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits.softmax(1)

    results = {}

    for scores in logits:
        for i, score in enumerate(scores):
            results[model.config.id2label[i]] = score.item()

    return results


demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(label="Input image", type="pil"),
    outputs="label",
)
demo.launch()