Upload app.py

app.py

@@ -3,7 +3,7 @@ AlexNet — 허깅페이스 Spaces 데모
 논문: Krizhevsky, Sutskever, Hinton (NeurIPS 2012)
 
 실행: Spaces에서 자동 실행 (app.py 이름 필수)
-로컬: pip install gradio torch pillow
+로컬: pip install gradio torch pillow torchvision
       python app.py
 """
 
@@ -23,19 +23,24 @@ class ConvBlock(nn.Module):
     """
     groups 인자로 논문의 GPU 분할 전략을 제어하는 범용 블록.
 
-    groups=1 : cross-GPU (전체 채널 연결) — Conv1·Conv3·FC
-    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).
-        논문도 실제로 Conv1은 3채널 입력 전체를 받아 96채널로 변환한 뒤
-        Conv2부터 GPU별로 채널을 분리함.
+        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 (정상)
     """
-    def __init__(self, in_ch, out_ch, kernel_size, stride=1, padding=0,
-                 groups=1, use_lrn=False, use_pool=False):
+    def __init__(self, in_ch, out_ch, kernel_size,
+                 stride=1, padding=0, groups=1,
+                 use_lrn=False, use_pool=False):
         super().__init__()
-        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size,
-                              stride=stride, padding=padding, groups=groups)
+        self.conv = nn.Conv2d(
+            in_ch, out_ch, kernel_size,
+            stride=stride, padding=padding, groups=groups,
+        )
         self.relu = nn.ReLU(inplace=True)
         self.lrn  = nn.LocalResponseNorm(5, alpha=1e-4, beta=0.75, k=2) if use_lrn  else None
         self.pool = nn.MaxPool2d(kernel_size=3, stride=2)                if use_pool else None
@@ -51,30 +56,32 @@ class AlexNet(nn.Module):
     """
     논문 Figure 2 완전 재현.
 
-    groups 전략:
-      Conv1        groups=1  (in=3, RGB는 2로 나눌 수 없음)
-      Conv2·4·5    groups=2  (parallel — GPU 분할)
-      Conv3·FC     groups=1  (cross-GPU — 전체 채널 연결)
-
-    레이어별 shape:
-      입력          (B,   3, 224, 224)
-      conv1 + pool  (B,  96,  27,  27)
-      conv2 + pool  (B, 256,  13,  13)
-      conv3         (B, 384,  13,  13)  ← cross-GPU
-      conv4         (B, 384,  13,  13)
-      conv5 + pool  (B, 256,   6,   6)
-      FC1·2·3       (B, 4096) → (B, 4096) → (B, num_labels)
+    레이어별 출력 shape:
+      입력              (B,   3, 224, 224)
+      Conv1 -> Pool1    (B,  96,  55,  55) -> (B,  96, 27, 27)
+      Conv2 -> Pool2    (B, 256,  27,  27) -> (B, 256, 13, 13)
+      Conv3             (B, 384,  13,  13)  <- cross-GPU
+      Conv4             (B, 384,  13,  13)
+      Conv5 -> Pool5    (B, 256,  13,  13) -> (B, 256,  6,  6)
+      Flatten           (B, 9216)
+      FC1->FC2->FC3     (B, 4096) -> (B, 4096) -> (B, 1000)
     """
     def __init__(self, num_labels=1000, dropout=0.5):
         super().__init__()
-        # Conv1: in=3(RGB) → groups=1 필수 (3은 2로 나눌 수 없음)
-        self.conv1 = ConvBlock(  3,  96, 11, stride=4, groups=1, use_lrn=True,  use_pool=True)
-        # Conv2·4·5: in_ch가 짝수 → groups=2 로 GPU 분할 재현
-        self.conv2 = ConvBlock( 96, 256,  5, padding=2, groups=2, use_lrn=True,  use_pool=True)
-        # Conv3: cross-GPU
-        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)
+
+        # 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.classifier = nn.Sequential(
             nn.Dropout(p=dropout),
             nn.Linear(256 * 6 * 6, 4096),
@@ -115,26 +122,22 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = AlexNet(num_labels=1000, dropout=0.5).to(DEVICE)
 model.eval()
 
-# config.json 에서 id2label 읽기 (없으면 인덱스로 표시)
 try:
     with open("config.json") as f:
         cfg = json.load(f)
-    ID2LABEL = cfg.get("id2label", {})
-    ID2LABEL = {int(k): v for k, v in ID2LABEL.items()}
+    ID2LABEL = {int(k): v for k, v in cfg.get("id2label", {}).items()}
 except Exception:
     ID2LABEL = {}
 
 
 # ──────────────────────────────────────────────────────────────
 # 3. 전처리 파이프라인
-#    논문 2절: 256×256 다운샘플 → 224×224 center crop → 픽셀 평균 차감
 # ──────────────────────────────────────────────────────────────
 
 TRANSFORM = T.Compose([
     T.Resize(256),
     T.CenterCrop(224),
     T.ToTensor(),
-    # ImageNet 픽셀 평균 차감 (논문 2절: "subtracting the mean activity")
     T.Normalize(mean=[0.485, 0.456, 0.406],
                 std=[0.229, 0.224, 0.225]),
 ])
@@ -145,26 +148,13 @@ TRANSFORM = T.Compose([
 # ──────────────────────────────────────────────────────────────
 
 def predict(image: Image.Image) -> dict:
-    """
-    PIL 이미지를 받아 Top-5 클래스 확률을 반환합니다.
-
-    Args:
-        image: Gradio가 넘겨주는 PIL.Image 객체
-
-    Returns:
-        {클래스명: 확률} 딕셔너리 — Gradio Label 컴포넌트용
-    """
     if image is None:
         return {}
-
-    tensor = TRANSFORM(image).unsqueeze(0).to(DEVICE)  # (1, 3, 224, 224)
-
+    tensor = TRANSFORM(image).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
-        logits = model(tensor)                          # (1, 1000)
-
-    probs = torch.softmax(logits, dim=-1)[0]            # (1000,)
+        logits = model(tensor)
+    probs = torch.softmax(logits, dim=-1)[0]
     top5_probs, top5_idx = probs.topk(5)
-
     return {
         ID2LABEL.get(idx.item(), f"class_{idx.item()}"): round(prob.item(), 4)
         for prob, idx in zip(top5_probs, top5_idx)
@@ -193,17 +183,17 @@ with gr.Blocks(title="AlexNet — 논문 재현") as demo:
 
     with gr.Accordion("모델 구조 (논문 Figure 2)", open=False):
         gr.Markdown("""
-        | 레이어 | 출력 shape       | 특이사항                          |
-        |--------|-----------------|----------------------------------|
-        | Conv1  | (B, 96, 27, 27)  | 11×11, stride 4, LRN, MaxPool, groups=2 |
-        | Conv2  | (B, 256, 13, 13) | 5×5, LRN, MaxPool, groups=2      |
-        | Conv3  | (B, 384, 13, 13) | 3×3, **cross-GPU** (groups=1)    |
-        | Conv4  | (B, 384, 13, 13) | 3×3, groups=2                    |
-        | Conv5  | (B, 256, 6, 6)   | 3×3, MaxPool, groups=2           |
-        | FC1·2  | (B, 4096)        | Dropout 0.5                      |
-        | FC3    | (B, 1000)        | 출력층                            |
-
-        총 파라미터: 약 **6,000만 개**
+        | 레이어 | 출력 shape      | padding | groups | 비고 |
+        |--------|-----------------|---------|--------|------|
+        | 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 정상.
         """)
 
     run_btn.click(fn=predict, inputs=image_input, outputs=label_output)