Upload app.py

app.py

@@ -19,13 +19,23 @@ import torchvision.transforms as T
 # 1. 모델 정의
 # ──────────────────────────────────────────────────────────────
 
-class ParallelConvBlock(nn.Module):
-    """Conv1·2·4·5: groups=2 로 논문의 GPU 분할 구조 재현."""
+class ConvBlock(nn.Module):
+    """
+    groups 인자로 논문의 GPU 분할 전략을 제어하는 범용 블록.
+
+    groups=1 : cross-GPU (전체 채널 연결) — Conv1·Conv3·FC
+    groups=2 : parallel  (채널을 반씩 독립 연산) — Conv2·Conv4·Conv5
+
+    Conv1이 groups=1인 이유:
+        in_channels=3 (RGB)는 groups=2로 나눌 수 없음 (3 % 2 != 0).
+        논문도 실제로 Conv1은 3채널 입력 전체를 받아 96채널로 변환한 뒤
+        Conv2부터 GPU별로 채널을 분리함.
+    """
     def __init__(self, in_ch, out_ch, kernel_size, stride=1, padding=0,
-                 use_lrn=False, use_pool=False):
+                 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=2)
+                              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
@@ -37,21 +47,15 @@ class ParallelConvBlock(nn.Module):
         return x
 
 
-class CrossConvBlock(nn.Module):
-    """Conv3: groups=1 로 cross-GPU 전체 채널 연결."""
-    def __init__(self, in_ch, out_ch, kernel_size, padding=0):
-        super().__init__()
-        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=padding, groups=1)
-        self.relu = nn.ReLU(inplace=True)
-
-    def forward(self, x):
-        return self.relu(self.conv(x))
-
-
 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)
@@ -63,11 +67,14 @@ class AlexNet(nn.Module):
     """
     def __init__(self, num_labels=1000, dropout=0.5):
         super().__init__()
-        self.conv1 = ParallelConvBlock(  3,  96, 11, stride=4, use_lrn=True,  use_pool=True)
-        self.conv2 = ParallelConvBlock( 96, 256,  5, padding=2, use_lrn=True,  use_pool=True)
-        self.conv3 = CrossConvBlock   (256, 384,  3, padding=1)
-        self.conv4 = ParallelConvBlock(384, 384,  3, padding=1)
-        self.conv5 = ParallelConvBlock(384, 256,  3, padding=1, use_pool=True)
+        # 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)
         self.classifier = nn.Sequential(
             nn.Dropout(p=dropout),
             nn.Linear(256 * 6 * 6, 4096),