train_split_tower_5scale.py

"""Split tower: FCOS-style detection head on cofiber features.

Separate cls/reg towers (std 3x3 + depthwise residual). P3 stride-8 via
transposed conv, top-down lateral connections. Cofiber decomposition
replaces the FPN. 2D sinusoidal positional encoding concatenated to the
backbone patch features before the stem.

Training: autocast bf16 forward, fp32 master params + fp32 moments, CUDA
graph capture of the forward+loss+backward+optimizer step, shard prefetch
on a dedicated copy stream, precomputed per-image FCOS targets stored in
the feature cache shards.
"""
import argparse
import json
import math
import os
import sys
import threading
import time
from queue import Queue

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

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

from target_cache import (
    precompute_targets_for_image, make_locations,
    STRIDES, SIZE_RANGES,
)
from cuda_graph_trainer import CudaGraphTrainStep

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 = int(os.environ.get("ARENA_RESOLUTION", "640"))
H = RESOLUTION // 16
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):
    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):
    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):
    layers = [ConvGNBlock(hidden) for _ in range(n_std)]
    layers += [DWResBlock(hidden) for _ in range(n_dw)]
    return nn.Sequential(*layers)


def make_sin_pos_emb(H, W, dim, device):
    assert dim % 4 == 0, "pos emb dim must be divisible by 4"
    d = dim // 4
    ys = torch.arange(H, device=device, dtype=torch.float32)
    xs = torch.arange(W, device=device, dtype=torch.float32)
    omega = torch.exp(torch.arange(d, device=device, dtype=torch.float32)
                      * -(math.log(10000.0) / d))
    pe_y = torch.zeros(H, d * 2, device=device)
    pe_y[:, 0::2] = torch.sin(ys[:, None] * omega[None, :])
    pe_y[:, 1::2] = torch.cos(ys[:, None] * omega[None, :])
    pe_x = torch.zeros(W, d * 2, device=device)
    pe_x[:, 0::2] = torch.sin(xs[:, None] * omega[None, :])
    pe_x[:, 1::2] = torch.cos(xs[:, None] * omega[None, :])
    pos = torch.zeros(dim, H, W, device=device)
    pos[:d*2] = pe_y.permute(1, 0)[:, :, None].expand(-1, H, W)
    pos[d*2:] = pe_x.permute(1, 0)[None, :, :].expand(H, -1, W).permute(1, 0, 2)
    return pos.unsqueeze(0)


class SplitTowerHead(nn.Module):
    def __init__(self, feat_dim=768, hidden=192, n_std_layers=5, n_dw_layers=4, n_scales=4,
                 pos_emb_dim=64, text_embed_path=None):
        super().__init__()
        self.n_scales = n_scales
        self.pos_emb_dim = pos_emb_dim
        n_total = n_scales + 1

        input_dim = feat_dim + pos_emb_dim
        self.scale_norms = nn.ModuleList([nn.GroupNorm(1, input_dim) for _ in range(n_scales)])
        self.stem = nn.Conv2d(input_dim, hidden, 1)
        self.stem_act = nn.GELU()
        self.p3_upsample = nn.ConvTranspose2d(hidden, hidden, 2, stride=2)
        self.p3_norm = nn.GroupNorm(min(32, hidden), hidden)
        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)])
        self.cls_tower = make_tower(hidden, n_std_layers, n_dw_layers)
        self.reg_tower = make_tower(hidden, n_std_layers, n_dw_layers)

        # Text-aligned cls head: learnable projection into the text embedding
        # space; cosine similarity against frozen CLIP class embeddings; learned
        # temperature (logit_scale) and per-class bias.
        if text_embed_path is None:
            text_embed_path = os.environ.get("COCO_TEXT_EMBED_PATH")
        assert text_embed_path and os.path.isfile(text_embed_path), \
            f"text_embed_path missing: {text_embed_path}"
        blob = torch.load(text_embed_path, map_location="cpu", weights_only=False)
        text_embed = blob["embeddings"].float()  # (C, D), already L2-normed
        assert text_embed.shape[0] == NUM_CLASSES
        self.text_embed_dim = text_embed.shape[1]
        self.register_buffer("text_embed", text_embed)
        self.cls_project = nn.Linear(hidden, self.text_embed_dim, bias=False)
        self.logit_scale = nn.Parameter(torch.tensor(math.log(1.0 / 0.07)))
        self.cls_bias = nn.Parameter(torch.full((NUM_CLASSES,), -math.log(99)))

        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))

    def forward(self, spatial):
        B, C, H_, W_ = spatial.shape
        pos = make_sin_pos_emb(H_, W_, self.pos_emb_dim, spatial.device).expand(B, -1, -1, -1)
        spatial = torch.cat([spatial, pos], dim=1)
        cofibers = cofiber_decompose(spatial, self.n_scales)
        cls_l, reg_l, ctr_l = [], [], []

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

        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))

        p3 = self.p3_norm(self.p3_upsample(scale_features[0]))
        all_features = [p3] + scale_features

        for i, x in enumerate(all_features):
            cls_feat = self.cls_tower(x)
            reg_feat = self.reg_tower(x)

            # Text-aligned classification: project feature to text embed dim,
            # cosine similarity against frozen class prototypes, learned temp.
            B_, _, Hi, Wi = cls_feat.shape
            f = cls_feat.permute(0, 2, 3, 1).reshape(-1, cls_feat.shape[1])
            f_proj = self.cls_project(f)
            f_norm = F.normalize(f_proj, p=2, dim=-1)
            # text_embed is already L2-normed; compute cosine sim
            logits = f_norm @ self.text_embed.t()  # (B*H*W, C)
            cls = (logits * self.logit_scale.exp() + self.cls_bias).reshape(
                B_, Hi, Wi, NUM_CLASSES).permute(0, 3, 1, 2)

            reg_raw = (self.reg_pred(reg_feat) * self.scale_params[i]).clamp(-10, 10)
            reg = reg_raw.exp()
            ctr = self.ctr_pred(reg_feat)
            cls_l.append(cls); reg_l.append(reg); ctr_l.append(ctr)
        return cls_l, reg_l, ctr_l


class _NullCtx:
    def __enter__(self): return self
    def __exit__(self, *a): pass


def load_shard_to_gpu(shard_path, copy_stream=None):
    """Load an amended shard, pre-stack per-image tensors into pinned CPU
    blocks, async H2D to GPU. spatial as bf16 (forward runs under autocast);
    targets as fp32 for loss precision."""
    shard = torch.load(shard_path, map_location="cpu", weights_only=False)
    n = len(shard)
    spatial = torch.stack([s["spatial"] for s in shard]).pin_memory()
    tgt_cls = torch.stack([s["tgt_cls"] for s in shard]).pin_memory()
    tgt_reg = torch.stack([s["tgt_reg"] for s in shard]).pin_memory()
    tgt_ctr = torch.stack([s["tgt_ctr"] for s in shard]).pin_memory()
    ctx = torch.cuda.stream(copy_stream) if copy_stream is not None else _NullCtx()
    with ctx:
        result = {
            "spatial": spatial.cuda(non_blocking=True).to(torch.bfloat16),
            "tgt_cls": tgt_cls.cuda(non_blocking=True).long(),
            "tgt_reg": tgt_reg.cuda(non_blocking=True).float(),
            "tgt_ctr": tgt_ctr.cuda(non_blocking=True).float(),
            "n": n,
        }
    return result


def shard_prefetcher(shard_paths, shard_order, queue, copy_stream):
    """Worker loading shards onto `copy_stream` and pushing them to `queue`
    so the H2D overlaps with main-thread compute. Pushes None at end."""
    for idx in shard_order:
        shard_gpu = load_shard_to_gpu(shard_paths[idx], copy_stream=copy_stream)
        copy_stream.synchronize()
        queue.put(shard_gpu)
    queue.put(None)


def amend_shard_with_targets(shard_path, device="cuda"):
    """Add precomputed FCOS targets to a shard in place. Idempotent."""
    shard = torch.load(shard_path, map_location="cpu", weights_only=False)
    if "tgt_cls" in shard[0]:
        return shard
    feat_sizes = [(H * 2, H * 2), (H, H), (H // 2, H // 2),
                  (H // 4, H // 4), (H // 8, H // 8)]
    locs_per_level = make_locations(feat_sizes, STRIDES, torch.device(device))
    all_locs = torch.cat(locs_per_level, 0)
    n_per_level = [loc.shape[0] for loc in locs_per_level]
    level_ranges = []
    cumsum = 0
    for i, n in enumerate(n_per_level):
        lo, hi = SIZE_RANGES[i]
        level_ranges.append((cumsum, cumsum + n, STRIDES[i], lo, hi))
        cumsum += n
    for entry in shard:
        boxes = entry["boxes"].to(device).float()
        labels = entry["labels"].to(device).long()
        tcls, treg, tctr = precompute_targets_for_image(
            boxes, labels, all_locs, level_ranges, device)
        entry["tgt_cls"] = tcls.to(torch.int16).cpu()
        entry["tgt_reg"] = treg.to(torch.float16).cpu()
        entry["tgt_ctr"] = tctr.to(torch.float16).cpu()
    torch.save(shard, shard_path)
    return shard


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hidden", type=int, default=192)
    parser.add_argument("--std-layers", type=int, default=5)
    parser.add_argument("--dw-layers", type=int, default=4)
    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("--no-graph", action="store_true",
                        help="Disable CUDA Graphs (eager bf16 fallback)")
    parser.add_argument("--no-precompute", action="store_true",
                        help="Skip shard target-amend pass")
    parser.add_argument("--resume", type=str, default=None,
                        help="Resume from a checkpoint path")
    parser.add_argument("--no-aug", action="store_true",
                        help="Disable feature-space augmentations")
    parser.add_argument("--mixup-alpha", type=float, default=0.2,
                        help="Beta distribution alpha for mixup (0 = disabled)")
    args = parser.parse_args()

    head = SplitTowerHead(hidden=args.hidden, n_std_layers=args.std_layers,
                        n_dw_layers=args.dw_layers, n_scales=4).cuda()
    n_params = sum(p.numel() for p in head.parameters())
    print("=" * 60)
    print(f"Split tower 5-scale: {args.hidden} hidden, "
          f"{args.std_layers} std + {args.dw_layers} dw layers per tower")
    print(f"  {n_params:,} params")
    print(f"  CUDA Graphs: {'enabled' if not args.no_graph else 'disabled'}")
    print("=" * 60, flush=True)

    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,
                                  capturable=not args.no_graph)
    # Persistent LR tensor; filled in place before each replay. LambdaLR
    # replaces `group['lr']` by assignment which orphans the tensor the
    # captured optimizer kernel points to — LR would freeze at capture time.
    lr_tensor = torch.tensor(args.lr, device="cuda", dtype=torch.float32)
    for g in optimizer.param_groups:
        g['lr'] = lr_tensor

    def lr_at(step):
        if step < warmup:
            return args.lr * step / max(warmup, 1)
        progress = (step - warmup) / max(total_steps - warmup, 1)
        return args.lr * 0.5 * (1 + math.cos(math.pi * progress))

    feat_sizes = [(H * 2, H * 2), (H, H), (H // 2, H // 2),
                  (H // 4, H // 4), (H // 8, H // 8)]
    locs_per_level = make_locations(feat_sizes, STRIDES, torch.device("cuda"))
    all_locs = torch.cat(locs_per_level, 0)
    n_total = all_locs.shape[0]
    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, "
          f"{args.epochs} epochs", flush=True)

    # Single marker in manifest avoids walking every shard to re-check the
    # amended flag on re-runs. Amend pass only runs if the marker is missing.
    manifest_path = os.path.join(CACHE_DIR, "manifest.json")
    if not args.no_precompute and not manifest.get("targets_amended"):
        print("\nAmending shards with precomputed FCOS targets (one-time cost)...", flush=True)
        t0 = time.time()
        for i, sp in enumerate(shard_paths):
            shard = amend_shard_with_targets(sp)
            del shard
            if (i + 1) % 10 == 0:
                print(f"  shard {i+1}/{n_shards} ({(time.time()-t0)/(i+1):.1f}s/shard)", flush=True)
        manifest["targets_amended"] = True
        with open(manifest_path, "w") as f:
            json.dump(manifest, f, indent=2)
        print(f"  done in {(time.time()-t0)/60:.1f} min", flush=True)

    # Resume support: reload head + optimizer state, skip ahead.
    start_step = 0
    if args.resume:
        ckpt = torch.load(args.resume, map_location="cuda", weights_only=False)
        head.load_state_dict(ckpt["head"])
        if "optimizer" in ckpt:
            optimizer.load_state_dict(ckpt["optimizer"])
        start_step = ckpt.get("step", 0)
        if start_step < 0: start_step = 0
        print(f"  Resumed from step {start_step}", flush=True)

    graph_step = CudaGraphTrainStep(head, optimizer, batch_size=args.batch_size,
                                    max_boxes=64, all_locs=all_locs)

    # Precompute a flip permutation for the flat target layout.
    # Each level's (H, W) grid gets flipped on W; the permutation maps
    # original flat index → flipped flat index across all 5 levels.
    flip_perm_parts = []
    for h, w in feat_sizes:
        grid = torch.arange(h * w, device="cuda").reshape(h, w)
        flip_perm_parts.append(grid.flip(-1).flatten())
    flip_perm = torch.cat(flip_perm_parts)

    print(f"\n  Training...\n", flush=True)
    head.train()
    global_step = start_step
    t0 = time.time()

    copy_stream = torch.cuda.Stream()

    for epoch in range(args.epochs):
        if global_step >= total_steps: break
        shard_order = torch.randperm(n_shards).tolist()
        epoch_t0 = time.time()

        prefetch_q = Queue(maxsize=1)
        prefetch_thread = threading.Thread(
            target=shard_prefetcher,
            args=(shard_paths, shard_order, prefetch_q, copy_stream),
            daemon=True,
        )
        prefetch_thread.start()

        for shard_pos in range(len(shard_order)):
            if global_step >= total_steps: break
            shard_gpu = prefetch_q.get()
            if shard_gpu is None: break
            n = shard_gpu["n"]
            within = torch.randperm(n, device="cuda")
            for batch_start in range(0, n, args.batch_size):
                if global_step >= total_steps: break
                batch_idx = within[batch_start:batch_start + args.batch_size]
                if batch_idx.shape[0] < args.batch_size: continue

                lr_tensor.fill_(lr_at(global_step))

                torch.index_select(shard_gpu["spatial"], 0, batch_idx, out=graph_step.buf_spatial)
                torch.index_select(shard_gpu["tgt_cls"], 0, batch_idx, out=graph_step.buf_tgt_cls)
                torch.index_select(shard_gpu["tgt_reg"], 0, batch_idx, out=graph_step.buf_tgt_reg)
                torch.index_select(shard_gpu["tgt_ctr"], 0, batch_idx, out=graph_step.buf_tgt_ctr)

                if not args.no_aug:
                    # Horizontal flip: 50% probability, applied per-batch.
                    # All ops write into the existing persistent buffers so the
                    # captured graph sees the augmented data on the next replay.
                    if torch.rand(1).item() < 0.5:
                        graph_step.buf_spatial.copy_(graph_step.buf_spatial.flip(-1))
                        graph_step.buf_tgt_cls.copy_(graph_step.buf_tgt_cls[:, flip_perm])
                        graph_step.buf_tgt_ctr.copy_(graph_step.buf_tgt_ctr[:, flip_perm])
                        flipped_reg = graph_step.buf_tgt_reg[:, flip_perm].clone()
                        flipped_reg[..., [0, 2]] = flipped_reg[..., [2, 0]]  # swap l↔r
                        graph_step.buf_tgt_reg.copy_(flipped_reg)

                    # Mixup: blend current batch with a random second batch from
                    # the same shard. Targets come from the dominant image (the
                    # one with higher lambda weight).
                    if args.mixup_alpha > 0 and n >= 2 * args.batch_size:
                        lam = torch.distributions.Beta(args.mixup_alpha, args.mixup_alpha).sample().item()
                        if lam < 0.5: lam = 1.0 - lam  # ensure lam >= 0.5
                        mix_idx = within[torch.randperm(n, device="cuda")[:args.batch_size]]
                        mix_sp = shard_gpu["spatial"][mix_idx]
                        graph_step.buf_spatial.mul_(lam).add_(mix_sp, alpha=1.0 - lam)

                if global_step == start_step and not args.no_graph:
                    print("Capturing CUDA graph (first batch is slow)...", flush=True)
                    try:
                        graph_step.warmup_and_capture()
                        print("  graph captured", flush=True)
                    except Exception as e:
                        print(f"  graph capture failed ({type(e).__name__}); falling back to eager: {e}", flush=True)
                        graph_step.captured = False

                graph_step.run()
                global_step += 1

                if global_step % 100 == 0:
                    loss_val = graph_step.last_loss()
                    elapsed = time.time() - t0
                    print(f"  step {global_step}/{total_steps} (ep {epoch+1}) "
                          f"loss={loss_val:.4f} lr={lr_tensor.item():.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", "split_tower_5scale_textaligned")
                    os.makedirs(out_dir, exist_ok=True)
                    ckpt_path = os.path.join(out_dir, f"checkpoint_step{global_step}.pth")
                    sd = head.state_dict()
                    torch.save({"head": {k: v.float() for k, v in sd.items()},
                                "optimizer": optimizer.state_dict(),
                                "step": global_step}, ckpt_path)
            del shard_gpu
            torch.cuda.empty_cache()

        while True:
            try:
                leftover = prefetch_q.get_nowait()
                if leftover is not None:
                    del leftover
            except Exception:
                break
        prefetch_thread.join(timeout=30)
        torch.cuda.empty_cache()
        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", "split_tower_5scale_textaligned")
    os.makedirs(out_dir, exist_ok=True)
    sd = head.state_dict()
    torch.save({"head": {k: v.float() for k, v in sd.items()}, "step": -1},
               os.path.join(out_dir, f"split_tower_5scale_{args.hidden}h_{args.std_layers}std_{args.dw_layers}dw_{args.epochs}ep.pth"))
    print(f"Done. Total time: {(time.time()-t0)/60:.1f} min")


if __name__ == "__main__":
    main()