add D3 = ResnetGenerator9Plain (no dilation, no wing ResBlock, bilinear upsample) for 1024^2 + DAB loss

scripts/src_ds_her2_thunder_v2.py

@@ -774,6 +774,78 @@ class ResnetGenerator9Dilated(nn.Module):
         return self.head(x)
 
 
+class ResnetGenerator9Plain(nn.Module):
+    """Arch D3: 1024^2 CycleGAN-style Resnet 9-block generator, NO dilation.
+
+    Standard Pyramid-Pix2Pix / CycleGAN template adapted to our stack:
+    - GroupNorm + GELU + zero pad throughout (NOT InstanceNorm + ReLU + reflection).
+    - Bilinear x2 + Conv3x3 for upsampling (NOT ConvTranspose, avoids checkerboard).
+    - Wing blocks are lean (Conv3x3+GN+GELU only, NO ResBlock) so we fit 1024^2 at
+      bs=8 on A100 80GB. All the ResBlocks live in the bottleneck.
+    - Bottleneck: 9 x ResBlock(base_ch*4), dilation=1 (plain), same class as A3/C.
+
+    Topology:
+        stem: Conv7x7(3->c1) + GN + GELU              (1024^2)
+        down1: Conv3x3(c1->c2, s=2) + GN + GELU        (512^2)
+        down2: Conv3x3(c2->c3, s=2) + GN + GELU        (256^2)
+        bottleneck: 9 x ResBlock(c3, dilation=1)       (256^2)
+        up1: Bilinear x2 + Conv3x3(c3->c2) + GN + GELU (512^2)
+        up2: Bilinear x2 + Conv3x3(c2->c1) + GN + GELU (1024^2)
+        head: Conv7x7(c1->3) + Tanh                    (1024^2)
+
+    ~10.4M params at base_ch=64. Works at any resolution divisible by 4.
+    """
+
+    def __init__(self, base_ch: int = 64, groups: int = 8, n_bottleneck: int = 9):
+        super().__init__()
+        c1, c2, c3 = base_ch, base_ch * 2, base_ch * 4
+        g1 = min(groups, c1); g2 = min(groups, c2); g3 = min(groups, c3)
+        self.stem = nn.Sequential(
+            nn.Conv2d(3, c1, 7, padding=3),
+            nn.GroupNorm(g1, c1),
+            nn.GELU(),
+        )
+        self.down1 = nn.Sequential(
+            nn.Conv2d(c1, c2, 3, stride=2, padding=1),
+            nn.GroupNorm(g2, c2),
+            nn.GELU(),
+        )
+        self.down2 = nn.Sequential(
+            nn.Conv2d(c2, c3, 3, stride=2, padding=1),
+            nn.GroupNorm(g3, c3),
+            nn.GELU(),
+        )
+        # Plain 9-block bottleneck (dilation=1), same ResBlock class as A3/C.
+        self.bottleneck = nn.Sequential(
+            *[ResBlock(c3, groups=groups) for _ in range(n_bottleneck)]
+        )
+        self.up1 = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+            nn.Conv2d(c3, c2, 3, padding=1),
+            nn.GroupNorm(g2, c2),
+            nn.GELU(),
+        )
+        self.up2 = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+            nn.Conv2d(c2, c1, 3, padding=1),
+            nn.GroupNorm(g1, c1),
+            nn.GELU(),
+        )
+        self.head = nn.Sequential(
+            nn.Conv2d(c1, 3, 7, padding=3),
+            nn.Tanh(),
+        )
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.down1(x)
+        x = self.down2(x)
+        x = self.bottleneck(x)
+        x = self.up1(x)
+        x = self.up2(x)
+        return self.head(x)
+
+
 class SkipFreeAE(nn.Module):
     """Encoder-decoder with pluggable downsample block (strided or DWT).
     Used for Exp 2 (arch A vs B) at ds_factor=1.
@@ -1190,6 +1262,8 @@ def train_one(arch: str, ds_factor: int, files_train, files_val,
     elif arch == "D2":
         G = ResnetGenerator9Dilated(base_ch=args.dilated_ch,
                                      dilations=tuple(args.d2_dilations)).to(device)
+    elif arch == "D3":
+        G = ResnetGenerator9Plain(base_ch=args.dilated_ch).to(device)
     elif arch == "C":
         G = WaveletJointAE(base_ch=args.wavelet_ch,
                             wavelet=args.wavelet).to(device)