Upload 2 files

app.py

@@ -2,17 +2,24 @@
 AlexNet — 허깅페이스 Spaces 데모
 논문: Krizhevsky, Sutskever, Hinton (NeurIPS 2012)
 
+변경사항:
+  - ImageNet id2label 자동 로드 (ViT config에서 가져옴)
+  - torchvision 사전학습 가중치 (FC 레이어) 로드
+  - 가중치 로드 실패 시 랜덤 초기화로 폴백
+
 실행: Spaces에서 자동 실행 (app.py 이름 필수)
-로컬: pip install gradio torch pillow torchvision
+로컬: pip install gradio torch pillow torchvision requests
       python app.py
 """
 
 import json
+import requests
 import torch
 import torch.nn as nn
+import torchvision.models as tv
+import torchvision.transforms as T
 import gradio as gr
 from PIL import Image
-import torchvision.transforms as T
 
 
 # ──────────────────────────────────────────────────────────────
@@ -22,16 +29,12 @@ import torchvision.transforms as T
 class ConvBlock(nn.Module):
     """
     groups 인자로 논문의 GPU 분할 전략을 제어하는 범용 블록.
+      groups=1: cross-GPU (전체 채널 연결) — Conv1·Conv3
+      groups=2: parallel  (채널 반씩 독립) — Conv2·Conv4·Conv5
 
-    groups=1 : cross-GPU — Conv1·Conv3 (전체 채널 연결)
-    groups=2 : parallel  — Conv2·Conv4·Conv5 (채널을 반씩 독립 연산)
-
-    Conv1이 groups=1인 이유:
-        in_channels=3(RGB)는 groups=2로 나눌 수 없음 (3 % 2 != 0).
-
-    padding 계산 근거 (Conv1):
-        padding=0 → (224-11)/4+1 = 54.25 → 내림 54 → Pool 후 26 → ... → FC 입력 6400 (오류)
-        padding=2 → (224-11+4)/4+1 = 55   → Pool 후 27 → ... → FC 입력 9216 (정상)
+    Conv1 padding=2 이유:
+      padding=0 → 출력 54(내림) → FC 입력 6400 오류
+      padding=2 → 출력 55(정확) → FC 입력 9216 정상
     """
     def __init__(self, in_ch, out_ch, kernel_size,
                  stride=1, padding=0, groups=1,
@@ -68,20 +71,11 @@ class AlexNet(nn.Module):
     """
     def __init__(self, num_labels=1000, dropout=0.5):
         super().__init__()
-
-        # Conv1: padding=2 필수 — 55x55 출력 보장
-        self.conv1 = ConvBlock(
-            3, 96, 11, stride=4, padding=2, groups=1,
-            use_lrn=True, use_pool=True,
-        )
-        self.conv2 = ConvBlock(
-            96, 256, 5, padding=2, groups=2,
-            use_lrn=True, use_pool=True,
-        )
-        self.conv3 = ConvBlock(256, 384, 3, padding=1, groups=1)
-        self.conv4 = ConvBlock(384, 384, 3, padding=1, groups=2)
-        self.conv5 = ConvBlock(384, 256, 3, padding=1, groups=2, use_pool=True)
-
+        self.conv1 = ConvBlock(  3,  96, 11, stride=4, padding=2, groups=1, use_lrn=True,  use_pool=True)
+        self.conv2 = ConvBlock( 96, 256,  5,           padding=2, groups=2, use_lrn=True,  use_pool=True)
+        self.conv3 = ConvBlock(256, 384,  3,           padding=1, groups=1)
+        self.conv4 = ConvBlock(384, 384,  3,           padding=1, groups=2)
+        self.conv5 = ConvBlock(384, 256,  3,           padding=1, groups=2, use_pool=True)
         self.classifier = nn.Sequential(
             nn.Dropout(p=dropout),
             nn.Linear(256 * 6 * 6, 4096),
@@ -114,24 +108,60 @@ class AlexNet(nn.Module):
 
 
 # ──────────────────────────────────────────────────────────────
-# 2. 모델 + 클래스 레이블 로드
+# 2. 모델 생성 + 사전학습 가중치 로드
 # ──────────────────────────────────────────────────────────────
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model  = AlexNet(num_labels=1000, dropout=0.5).to(DEVICE)
+
+WEIGHTS_LOADED = False
+try:
+    # torchvision AlexNet 사전학습 가중치에서 FC 레이어만 복사
+    # (Conv 레이어는 groups 구조가 달라 직접 로드 불가)
+    pretrained = tv.alexnet(weights=tv.AlexNet_Weights.DEFAULT)
+    model.classifier.load_state_dict(pretrained.classifier.state_dict())
+    WEIGHTS_LOADED = True
+    print("사전학습 가중치(FC) 로드 완료")
+except Exception as e:
+    print(f"가중치 로드 실패, 랜덤 초기화 유지: {e}")
 
-model = AlexNet(num_labels=1000, dropout=0.5).to(DEVICE)
 model.eval()
 
+
+# ──────────────────────────────────────────────────────────────
+# 3. ImageNet id2label 로드
+#    우선순위: config.json → ViT config(허깅페이스) → 인덱스 표시
+# ──────────────────────────────────────────────────────────────
+
+ID2LABEL = {}
+
+# 1순위: config.json
 try:
     with open("config.json") as f:
         cfg = json.load(f)
     ID2LABEL = {int(k): v for k, v in cfg.get("id2label", {}).items()}
+    if ID2LABEL:
+        print(f"config.json에서 {len(ID2LABEL)}개 클래스 로드")
 except Exception:
-    ID2LABEL = {}
+    pass
+
+# 2순위: 허깅페이스 ViT config (ImageNet 1000개 라벨 동일)
+if not ID2LABEL:
+    try:
+        resp = requests.get(
+            "https://huggingface.co/google/vit-base-patch16-224/raw/main/config.json",
+            timeout=10,
+        )
+        vit_cfg = resp.json()
+        ID2LABEL = {int(k): v for k, v in vit_cfg.get("id2label", {}).items()}
+        print(f"허깅페이스에서 {len(ID2LABEL)}개 클래스 로드")
+    except Exception as e:
+        print(f"클래스 이름 로드 실패: {e}")
 
 
 # ──────────────────────────────────────────────────────────────
-# 3. 전처리 파이프라인
+# 4. 전처리 파이프라인
+#    논문 2절: 256 리사이즈 → 224 center crop → 픽셀 평균 차감
 # ──────────────────────────────────────────────────────────────
 
 TRANSFORM = T.Compose([
@@ -144,7 +174,7 @@ TRANSFORM = T.Compose([
 
 
 # ──────────────────────────────────────────────────────────────
-# 4. 추론 함수
+# 5. 추론 함수
 # ──────────────────────────────────────────────────────────────
 
 def predict(image: Image.Image) -> dict:
@@ -162,16 +192,19 @@ def predict(image: Image.Image) -> dict:
 
 
 # ──────────────────────────────────────────────────────────────
-# 5. Gradio UI
+# 6. Gradio UI
 # ──────────────────────────────────────────────────────────────
 
+weight_status = "FC 사전학습 가중치 로드됨 (torchvision)" if WEIGHTS_LOADED else "랜덤 초기화 상태"
+label_status  = f"ImageNet {len(ID2LABEL)}개 클래스 이름 로드됨" if ID2LABEL else "클래스 이름 없음 (인덱스 표시)"
+
 with gr.Blocks(title="AlexNet — 논문 재현") as demo:
-    gr.Markdown("""
+    gr.Markdown(f"""
     ## AlexNet — 논문 완전 재현 데모
     **논문**: ImageNet Classification with Deep CNNs (Krizhevsky et al., NeurIPS 2012)
 
-    > 이미지를 업로드하면 Top-5 클래스를 예측합니다.  
-    > ※ 현재 모델은 랜덤 초기화 상태입니다. ImageNet 학습 가중치를 로드하면 실제 예측이 가능합니다.
+    - 가중치: {weight_status}
+    - 클래스: {label_status}
     """)
 
     with gr.Row():
@@ -185,15 +218,13 @@ with gr.Blocks(title="AlexNet — 논문 재현") as demo:
         gr.Markdown("""
         | 레이어 | 출력 shape      | padding | groups | 비고 |
         |--------|-----------------|---------|--------|------|
-        | Conv1  | (B, 96, 55→27)  | **2**   | 1      | 11x11 stride4, LRN, MaxPool |
+        | Conv1  | (B, 96, 55→27)  | 2       | 1      | 11x11 stride4, LRN, MaxPool |
         | Conv2  | (B, 256, 27→13) | 2       | 2      | 5x5, LRN, MaxPool |
         | Conv3  | (B, 384, 13)    | 1       | 1      | 3x3, cross-GPU |
         | Conv4  | (B, 384, 13)    | 1       | 2      | 3x3 |
         | Conv5  | (B, 256, 13→6)  | 1       | 2      | 3x3, MaxPool |
-        | FC1·2  | (B, 4096)       | —       | —      | Dropout 0.5 |
-        | FC3    | (B, 1000)       | —       | —      | 출력층 |
-
-        **수정 포인트**: Conv1 padding=0이면 출력 54(내림)→FC 입력 6400 오류. padding=2로 55→FC 입력 9216 정상.
+        | FC1·2  | (B, 4096)       | —       | —      | Dropout 0.5, 사전학습 가중치 |
+        | FC3    | (B, 1000)       | —       | —      | 출력층, 사전학습 가중치 |
         """)
 
     run_btn.click(fn=predict, inputs=image_input, outputs=label_output)

requirements.txt

@@ -2,3 +2,4 @@ torch>=2.0.0
 torchvision>=0.15.0
 gradio>=4.0.0
 pillow>=9.0.0
+requests>=2.28.0