"""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()