vit-cifar10-classifier

ViT Fake Image Detector

このモデルは AI生成画像検出用にファインチューニングされた Vision Transformer です。

モデル情報

• ベースモデル: timm/vit_tiny_patch16_224
• データセット: uoft-cs/cifar10
• タスク: 2クラス分類 (fake vs real)
• 最高検証精度: 70.61%

クラス

使用方法

import torch
import timm
from PIL import Image
from torchvision import transforms

# モデル読み込み
model = timm.create_model("vit_tiny_patch16_224", pretrained=False, num_classes=2)
state_dict = torch.hub.load_state_dict_from_url(
    "https://huggingface.co/shunya510hi/vit-cifar10-classifier/resolve/main/pytorch_model.bin",
    map_location="cpu"
)
model.load_state_dict(state_dict)
model.eval()

# 画像前処理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 推論
image = Image.open("test_image.jpg").convert("RGB")
input_tensor = transform(image).unsqueeze(0)

with torch.no_grad():
    output = model(input_tensor)
    probabilities = torch.softmax(output, dim=1)
    prediction = output.argmax(1).item()

class_names = ["fake", "real"]
print(f"Prediction: {class_names[prediction]}")
print(f"Confidence: {probabilities[0][prediction].item():.2%}")

学習設定

• Epochs: 3
• Batch Size: 32
• Learning Rate: 0.0001
• Weight Decay: 0.01
• 学習対象: 最後の Transformer ブロック + head のみ