scripts/train_split_tower.py

"""
FCOS-Lite: Slim FCOS-style head on cofiber features.

Separate cls and reg towers with standard 3x3 convolutions (full cross-channel mixing).
P3 stride-8 via transposed conv. Top-down lateral connections.
Cofiber decomposition replaces the heavy FPN.

Target: match FCOS (41.0 mAP at 16.14M) at ≤4M params.

Key differences from conv_deep:
  - Standard Conv2d(256, 256, 3) instead of depthwise (256× more params per layer but full mixing)
  - Separate cls and reg towers (FCOS-style)
  - Fewer blocks (4 per tower instead of 20 shared)
"""

import argparse
import json
import math
import os
import sys
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPT_DIR)

CACHE_DIR = os.environ.get("ARENA_CACHE_DIR")
COCO_ROOT = os.environ.get("ARENA_COCO_ROOT")
VAL_CACHE = os.environ.get("ARENA_VAL_CACHE")
RESOLUTION = 640
NUM_CLASSES = 80


def cofiber_decompose(f, n_scales):
    cofibers = []; residual = f
    for _ in range(n_scales - 1):
        omega = F.avg_pool2d(residual, 2)
        sigma_omega = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False)
        cofibers.append(residual - sigma_omega); residual = omega
    cofibers.append(residual); return cofibers


class ConvGNBlock(nn.Module):
    """Standard 3x3 conv + GroupNorm + GELU. Full cross-channel mixing."""
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
        self.norm = nn.GroupNorm(min(32, channels), channels)
        self.act = nn.GELU()

    def forward(self, x):
        return self.act(self.norm(self.conv(x)))


class DWResBlock(nn.Module):
    """Depthwise residual block: pointwise + GELU + DW 3x3 + GN + residual."""
    def __init__(self, channels):
        super().__init__()
        self.pw = nn.Conv2d(channels, channels, 1)
        self.act = nn.GELU()
        self.dw = nn.Conv2d(channels, channels, 3, padding=1, groups=channels)
        self.norm = nn.GroupNorm(min(32, channels), channels)

    def forward(self, x):
        return x + self.norm(self.dw(self.act(self.pw(x))))


def make_tower(hidden, n_std, n_dw):
    """Build a hybrid tower: standard 3x3 layers + depthwise residual blocks."""
    layers = []
    for _ in range(n_std):
        layers.append(ConvGNBlock(hidden))
    for _ in range(n_dw):
        layers.append(DWResBlock(hidden))
    return nn.Sequential(*layers)


class FCOSLiteHead(nn.Module):
    """Slim FCOS head on cofiber features with P3 + lateral + hybrid towers."""

    def __init__(self, feat_dim=768, hidden=256, n_std_layers=3, n_dw_layers=6, n_scales=3):
        super().__init__()
        self.n_scales = n_scales
        n_total = n_scales + 1  # +1 for P3

        # Per-scale input norms
        self.scale_norms = nn.ModuleList([nn.GroupNorm(1, feat_dim) for _ in range(n_scales)])

        # Stem: project to hidden channels
        self.stem = nn.Conv2d(feat_dim, hidden, 1)
        self.stem_act = nn.GELU()

        # P3 upsample (stride 16 -> stride 8)
        self.p3_upsample = nn.ConvTranspose2d(hidden, hidden, 2, stride=2)
        self.p3_norm = nn.GroupNorm(min(32, hidden), hidden)

        # Top-down lateral connections
        self.lateral_convs = nn.ModuleList([nn.Conv2d(hidden, hidden, 1) for _ in range(n_scales - 1)])
        self.lateral_norms = nn.ModuleList([nn.GroupNorm(min(32, hidden), hidden) for _ in range(n_scales - 1)])

        # Separate cls and reg towers: standard 3x3 + depthwise residual
        self.cls_tower = make_tower(hidden, n_std_layers, n_dw_layers)
        self.reg_tower = make_tower(hidden, n_std_layers, n_dw_layers)

        # Prediction heads
        self.cls_pred = nn.Conv2d(hidden, NUM_CLASSES, 1)
        self.reg_pred = nn.Conv2d(hidden, 4, 1)
        self.ctr_pred = nn.Conv2d(hidden, 1, 1)
        self.scale_params = nn.Parameter(torch.ones(n_total))

        # Initialize cls bias for focal loss
        nn.init.constant_(self.cls_pred.bias, -math.log(99))

    def forward(self, spatial):
        cofibers = cofiber_decompose(spatial, self.n_scales)
        cls_l, reg_l, ctr_l = [], [], []

        # Process each scale through stem
        scale_features = []
        for i, cof in enumerate(cofibers):
            x = self.stem_act(self.stem(self.scale_norms[i](cof)))
            scale_features.append(x)

        # Top-down lateral fusion (coarse -> fine)
        for i in range(len(scale_features) - 2, -1, -1):
            coarse_up = F.interpolate(scale_features[i + 1], size=scale_features[i].shape[2:],
                                      mode="bilinear", align_corners=False)
            scale_features[i] = self.lateral_norms[i](
                scale_features[i] + self.lateral_convs[i](coarse_up))

        # Create P3 from stride-16 features
        p3 = self.p3_norm(self.p3_upsample(scale_features[0]))
        all_features = [p3] + scale_features

        # Run cls and reg towers on each level, predict
        for i, x in enumerate(all_features):
            cls_feat = self.cls_tower(x)
            reg_feat = self.reg_tower(x)

            cls = self.cls_pred(cls_feat)
            reg_raw = (self.reg_pred(reg_feat) * self.scale_params[i]).clamp(-10, 10)
            reg = reg_raw.exp()
            ctr = self.ctr_pred(reg_feat)  # centerness from reg tower (FCOS convention)

            cls_l.append(cls); reg_l.append(reg); ctr_l.append(ctr)
        return cls_l, reg_l, ctr_l


def make_locations(feature_sizes, strides, device):
    locs = []
    for (h, w), s in zip(feature_sizes, strides):
        ys = (torch.arange(h, device=device, dtype=torch.float32) + 0.5) * s
        xs = (torch.arange(w, device=device, dtype=torch.float32) + 0.5) * s
        gy, gx = torch.meshgrid(ys, xs, indexing="ij")
        locs.append(torch.stack([gx.flatten(), gy.flatten()], -1))
    return locs


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hidden", type=int, default=224)
    parser.add_argument("--std-layers", type=int, default=3)
    parser.add_argument("--dw-layers", type=int, default=6)
    parser.add_argument("--epochs", type=int, default=8)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--resume", type=str, default=None)
    args = parser.parse_args()

    head = FCOSLiteHead(hidden=args.hidden, n_std_layers=args.std_layers, n_dw_layers=args.dw_layers).cuda()
    start_step = 0
    if args.resume:
        ckpt = torch.load(args.resume, map_location="cuda", weights_only=False)
        head.load_state_dict(ckpt["head"])
        start_step = ckpt["step"]
        print(f"Resumed from step {start_step}")
    n_params = sum(p.numel() for p in head.parameters())
    print("=" * 60)
    print(f"FCOS-Lite: {args.hidden} hidden, {args.std_layers} std + {args.dw_layers} dw layers per tower")
    print(f"  {n_params:,} params")
    print(f"  Separate cls/reg towers, standard 3x3 convs, P3 + lateral")
    print("=" * 60, flush=True)

    from cache_and_train_fast import compute_loss
    manifest = json.load(open(os.path.join(CACHE_DIR, "manifest.json")))
    n_shards = manifest["n_shards"]
    n_images = manifest["n_images"]
    steps_per_epoch = n_images // args.batch_size
    total_steps = steps_per_epoch * args.epochs
    warmup = int(total_steps * 0.03)

    optimizer = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda s:
        s / max(warmup, 1) if s < warmup else
        0.5 * (1 + math.cos(math.pi * (s - warmup) / max(total_steps - warmup, 1))))
    scaler = GradScaler()

    H = RESOLUTION // 16
    strides = [8, 16, 32, 64]
    locs = make_locations([(H*2,H*2),(H,H),(H//2,H//2),(H//4,H//4)], strides, torch.device("cuda"))
    shard_paths = [os.path.join(CACHE_DIR, f"shard_{i:04d}.pt") for i in range(n_shards)]

    print(f"  {n_images} images, batch {args.batch_size}, {total_steps} steps, {args.epochs} epochs")
    print(f"  fp16 mixed precision")
    print(f"  Training...\n", flush=True)

    head.train()
    global_step = start_step
    if start_step > 0:
        for _ in range(start_step):
            scheduler.step()
        print(f"  Scheduler advanced to step {start_step}", flush=True)
    t0 = time.time()

    for epoch in range(args.epochs):
        shard_order = torch.randperm(n_shards).tolist()
        epoch_t0 = time.time()

        for shard_idx in shard_order:
            if global_step >= total_steps: break
            shard = torch.load(shard_paths[shard_idx], map_location="cpu", weights_only=False)
            within = torch.randperm(len(shard)).tolist()

            for batch_start in range(0, len(shard), args.batch_size):
                if global_step >= total_steps: break
                batch_idx = within[batch_start:batch_start + args.batch_size]
                if len(batch_idx) < 2: continue

                spatial = torch.stack([shard[i]["spatial"] for i in batch_idx]).float().cuda()
                boxes = [shard[i]["boxes"].cuda() for i in batch_idx]
                labels = [shard[i]["labels"].cuda() for i in batch_idx]

                try:
                    with autocast():
                        cls_l, reg_l, ctr_l = head(spatial)
                        loss = compute_loss(cls_l, reg_l, ctr_l, locs, boxes, labels)

                    optimizer.zero_grad()
                    scaler.scale(loss).backward()
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(head.parameters(), 5.0)
                    scaler.step(optimizer)
                    scaler.update()
                    scheduler.step()

                    global_step += 1

                    if global_step % 100 == 0:
                        lr = scheduler.get_last_lr()[0]
                        elapsed = time.time() - t0
                        print(f"  step {global_step}/{total_steps} (ep {epoch+1}) "
                              f"loss={loss.item():.4f} lr={lr:.2e} "
                              f"{global_step/elapsed:.1f} it/s", flush=True)

                    if global_step % 4000 == 0:
                        out_dir = os.path.join(SCRIPT_DIR, "heads", "cofiber_threshold", "fcos_lite")
                        os.makedirs(out_dir, exist_ok=True)
                        ckpt = os.path.join(out_dir, f"checkpoint_step{global_step}.pth")
                        torch.save({"head": head.state_dict(), "step": global_step}, ckpt)

                except RuntimeError as e:
                    if "out of memory" in str(e):
                        torch.cuda.empty_cache()
                        optimizer.zero_grad()
                        global_step += 1
                        scheduler.step()
                        continue
                    raise

            del shard

        print(f"  Epoch {epoch+1}/{args.epochs} complete ({time.time()-epoch_t0:.0f}s)\n", flush=True)

    out_dir = os.path.join(SCRIPT_DIR, "heads", "cofiber_threshold", "fcos_lite")
    os.makedirs(out_dir, exist_ok=True)
    out = os.path.join(out_dir, f"fcos_lite_{args.hidden}h_{args.tower_layers}t_{args.epochs}ep.pth")
    torch.save(head.state_dict(), out)
    elapsed = time.time() - t0
    print(f"Saved: {out}")
    print(f"{n_params:,} params, {elapsed/60:.1f} minutes")


if __name__ == "__main__":
    main()