train_food.py

"""Food-101 데이터셋으로 MyResNet을 학습하는 스크립트.

torchvision의 ImageNet pretrained ResNet-18 가중치를 가져와서
MyResNet에 로드한 후 Food-101에 fine-tuning합니다.

사용법:
    python train_food.py

요구사항:
    pip install torch torchvision transformers datasets accelerate

학습 시간 (GPU 1장 기준):
    - Food-101 full (101종): 약 2~4시간 (10 epoch)
    - 빠른 테스트: epochs=3으로 줄이면 1시간 내
"""
import numpy as np
import torch
import torch.nn as nn
import torchvision.models as tv_models
from datasets import load_dataset
from torchvision.transforms import (
    Compose,
    Normalize,
    RandomCrop,
    RandomHorizontalFlip,
    Resize,
    ToTensor,
)
from transformers import DefaultDataCollator, Trainer, TrainingArguments

from configuration_myresnet import MyResNetConfig
from modeling_myresnet import MyResNetForImageClassification


# ============================================================
# 1) 데이터셋 로딩
# ============================================================
print("Food-101 데이터셋 로딩 중...")
dataset = load_dataset("food101")

# Food-101의 101개 클래스명 가져오기
food_classes = dataset["train"].features["label"].names
NUM_CLASSES = len(food_classes)
print(f"총 클래스 수: {NUM_CLASSES}")
print(f"학습 이미지: {len(dataset['train']):,}장")
print(f"검증 이미지: {len(dataset['validation']):,}장")


# ============================================================
# 2) 모델 준비 - ResNet-18 + Pretrained 가중치 로드
# ============================================================
config = MyResNetConfig(
    num_channels=3,
    num_labels=NUM_CLASSES,
    block_type="basic",
    layers=[2, 2, 2, 2],                # ResNet-18
    hidden_sizes=[64, 128, 256, 512],
    image_size=224,
    id2label={i: name for i, name in enumerate(food_classes)},
    label2id={name: i for i, name in enumerate(food_classes)},
)
model = MyResNetForImageClassification(config)


def load_pretrained_resnet18(model):
    """torchvision의 ImageNet pretrained ResNet-18 가중치를 MyResNet에 복사.

    두 모델의 레이어 이름이 다르므로 매핑해줍니다.
    마지막 FC 레이어(classifier)는 클래스 수가 다르므로 스킵.
    """
    print("ImageNet pretrained ResNet-18 가중치 로딩...")
    tv_resnet = tv_models.resnet18(weights=tv_models.ResNet18_Weights.IMAGENET1K_V1)
    tv_state = tv_resnet.state_dict()

    # torchvision -> MyResNet 레이어 이름 매핑
    # tv_resnet:  conv1, bn1, layer1~4, fc
    # MyResNet:   stem.0 (conv1), stem.1 (bn1), stage1~4, classifier
    mapping = {
        "conv1.": "stem.0.",
        "bn1.": "stem.1.",
        "layer1.": "stage1.",
        "layer2.": "stage2.",
        "layer3.": "stage3.",
        "layer4.": "stage4.",
    }

    # downsample은 shortcut으로 매핑 (각 stage의 첫 블록에만 있음)
    new_state = {}
    for k, v in tv_state.items():
        if k.startswith("fc."):
            continue  # FC 레이어는 스킵 (클래스 수가 다름)

        new_k = k
        for old, new in mapping.items():
            if new_k.startswith(old):
                new_k = new_k.replace(old, new, 1)
                break

        # layer1.0.downsample.0 -> stage1.0.shortcut.0
        new_k = new_k.replace(".downsample.", ".shortcut.")
        new_state[new_k] = v

    # strict=False로 부분 로드 (classifier는 학습 필요)
    missing, unexpected = model.load_state_dict(new_state, strict=False)
    print(f"로드 성공. 누락된 키 {len(missing)}개 (classifier 등), "
          f"예상치 못한 키 {len(unexpected)}개")
    return model


model = load_pretrained_resnet18(model)
print(f"모델 파라미터 수: {sum(p.numel() for p in model.parameters()):,}")


# ============================================================
# 3) 데이터 전처리
# ============================================================
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

train_transform = Compose([
    Resize(256),
    RandomCrop(224),
    RandomHorizontalFlip(),
    ToTensor(),
    Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])
eval_transform = Compose([
    Resize((224, 224)),
    ToTensor(),
    Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])


def preprocess_train(batch):
    batch["pixel_values"] = [
        train_transform(img.convert("RGB")) for img in batch["image"]
    ]
    batch["labels"] = batch["label"]
    return batch


def preprocess_eval(batch):
    batch["pixel_values"] = [
        eval_transform(img.convert("RGB")) for img in batch["image"]
    ]
    batch["labels"] = batch["label"]
    return batch


train_ds = dataset["train"].with_transform(preprocess_train)
eval_ds = dataset["validation"].with_transform(preprocess_eval)


# ============================================================
# 4) 평가 지표 (top-1 + top-5 accuracy)
# ============================================================
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    # Top-1 accuracy
    top1_preds = np.argmax(logits, axis=-1)
    top1_acc = (top1_preds == labels).mean()

    # Top-5 accuracy
    top5_preds = np.argsort(-logits, axis=-1)[:, :5]
    top5_correct = np.any(top5_preds == labels[:, None], axis=-1)
    top5_acc = top5_correct.mean()

    return {
        "accuracy": float(top1_acc),
        "top5_accuracy": float(top5_acc),
    }


# ============================================================
# 5) Trainer 설정
# ============================================================
training_args = TrainingArguments(
    output_dir="./my-resnet18-food101",
    num_train_epochs=10,                    # 10 epoch 권장 (pretrained이므로 적게)
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    learning_rate=1e-3,                     # pretrained이므로 더 작게 시작
    weight_decay=1e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
    eval_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    logging_steps=50,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    greater_is_better=True,
    fp16=torch.cuda.is_available(),
    remove_unused_columns=False,
    dataloader_num_workers=4,
    report_to="none",
    push_to_hub=False,                      # True로 바꾸면 자동 업로드
    # hub_model_id="your-username/my-resnet18-food101",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    data_collator=DefaultDataCollator(),
    compute_metrics=compute_metrics,
)


# ============================================================
# 6) 학습 실행
# ============================================================
if __name__ == "__main__":
    print("\n=== Food-101 학습 시작 ===")
    trainer.train()

    metrics = trainer.evaluate()
    print(f"\n최종 Top-1 정확도: {metrics['eval_accuracy']:.4f}")
    print(f"최종 Top-5 정확도: {metrics['eval_top5_accuracy']:.4f}")

    trainer.save_model("./my-resnet18-food101")
    print("모델 저장 완료: ./my-resnet18-food101")

    # 허브 업로드:
    # trainer.push_to_hub()