diff --git a/eomt/api.py b/eomt/api.py index b65f786..061f031 100644 --- a/eomt/api.py +++ b/eomt/api.py @@ -396,6 +396,7 @@ def save(self, path: str | Path) -> None: aux_heads=m.aux_specs, aux_head_arch=m.aux_head_arch, letterbox=bool(getattr(m, "preprocess_letterbox", False)), loss_weights=m.loss_weights, num_upscale_blocks=m.num_upscale_blocks, + fpn_scales=getattr(m, "fpn_scales", None), norm_mean=getattr(m, "pixel_mean", None), norm_std=getattr(m, "pixel_std", None), patch_size=getattr(m, "patch_size", None), compression=self._compression, diff --git a/eomt/engine/train.py b/eomt/engine/train.py index 7b073e3..f9484cf 100644 --- a/eomt/engine/train.py +++ b/eomt/engine/train.py @@ -30,7 +30,7 @@ from ..data import CocoDetection, CocoInstanceSeg, CocoValImages, collate_train from ..data.transforms import build_val_transform from ..ema import ModelEMA, _unwrap -from ..model import build_model, load_dinov2_backbone +from ..model import DEFAULT_FPN_SCALES, build_model, load_dinov2_backbone from ..plotting import plot_aux_per_class, plot_metrics_csv from ..device import resolve_device from ..serialization import load_raw, resolve_checkpoint, save_checkpoint, wrap_checkpoint @@ -306,6 +306,7 @@ def train( l1_weight: float = 5.0, giou_weight: float = 2.0, num_upscale_blocks: int | None = None, + fpn_scales: tuple[float, ...] | list[float] | None = DEFAULT_FPN_SCALES, logger: str = "none", resume: str | None = None, init_weights: str | None = None, @@ -319,6 +320,12 @@ def train( ``flip_prob=0``) while keeping the learned weights. The two are mutually exclusive. + ``fpn_scales`` defaults to the B1 multi-scale recipe (SimpleFPN + query + cross-attention, :data:`~eomt.model.DEFAULT_FPN_SCALES`); pass ``None`` for the + original single-scale model. Resuming ignores this and keeps the checkpoint's + architecture; warm-starting (``init_weights``) honors it, so fine-tuning a + single-scale checkpoint into a multi-scale one is the default warm-start path. + ``freeze_backbone_epochs`` trains the task head only for the first N epochs (DINOv2 frozen), then unfreezes — the LP-FT recipe. With a from-scratch head this protects the pretrained features from early noisy-gradient corruption and @@ -445,8 +452,17 @@ def train( if resume_ckpt is not None: loss_weights = normalize_loss_weights(resume_ckpt.get("loss_weights"), family=family) num_upscale_blocks = resume_ckpt.get("num_upscale_blocks", num_upscale_blocks) + # Resuming MUST keep the checkpoint's architecture (the saved FPN/cross-attn + # tensors must match the rebuilt modules); the CLI flag/default is ignored. + # An absent key means the checkpoint is single-scale — do NOT fall back to the + # (now multi-scale) default, or resuming an old run would build a mismatched head. + fpn_scales = resume_ckpt.get("fpn_scales") elif init_ckpt is not None and init_ckpt.get("num_upscale_blocks") is not None: num_upscale_blocks = int(init_ckpt["num_upscale_blocks"]) + # Warm-start (init_weights) deliberately keeps the CLI ``fpn_scales``: enabling + # multi-scale by fine-tuning a single-scale checkpoint is the intended workflow — + # the new FPN/cross-attn modules init random (shape-mismatched/missing tensors are + # skipped by the warm-start loader below) while the ViT + heads load 1:1. # --- data --- Dataset = CocoDetection if is_detect else CocoInstanceSeg @@ -549,6 +565,7 @@ def train( "aux_heads": {s.name: s.num_classes for s in aux_specs}, "loss_weights": loss_weights, "num_upscale_blocks": num_upscale_blocks, + "fpn_scales": list(fpn_scales) if fpn_scales else None, }, ) @@ -563,8 +580,11 @@ def train( aux_head_arch=aux_head_arch, loss_weights=loss_weights, num_upscale_blocks=num_upscale_blocks, + fpn_scales=fpn_scales, ).to(dev) print(f"[model] loss weights: {loss_weights}; upscale blocks: {model.num_upscale_blocks}") + if model.fpn_scales: + print(f"[model] multi-scale FPN enabled: scales={model.fpn_scales}") if aux_specs: print(f"[model] aux head arch: {model.aux_head_arch}") start_epoch, best_metric = 0, -1.0 @@ -677,6 +697,7 @@ def _save(path: Path, *, state_dict, with_trainer_state: bool): letterbox=letterbox, loss_weights=loss_weights, num_upscale_blocks=core.num_upscale_blocks, + fpn_scales=core.fpn_scales, norm_mean=getattr(core, "pixel_mean", None), norm_std=getattr(core, "pixel_std", None), patch_size=getattr(core, "patch_size", None), diff --git a/eomt/model.py b/eomt/model.py index 7c3d04b..1891840 100644 --- a/eomt/model.py +++ b/eomt/model.py @@ -273,6 +273,136 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.fc3(hidden_states).sigmoid() +#: Default ViTDet-style pyramid scales relative to the native stride-14 grid: +#: 2x (fine, the small-object win), 1x (native), 0.5x (coarse). Used when +#: ``fpn_scales`` is enabled without an explicit tuple, and by the serialization +#: inferer (the conv weights don't encode the scale, so this is the recovery prior). +DEFAULT_FPN_SCALES: tuple[float, ...] = (2.0, 1.0, 0.5) + + +def _sincos_2d(h: int, w: int, dim: int, device, dtype) -> torch.Tensor: + """Standard DETR 2D sine-cosine positional embedding, ``[h*w, dim]``. + + Computed on the fly per level (no fixed table) so the model still runs at any + input size. ``dim`` must be divisible by 4 (half the channels encode the row + coordinate, half the column; each half split into sin/cos). + """ + if dim % 4: + raise ValueError(f"_sincos_2d needs dim divisible by 4, got {dim}.") + quarter = dim // 4 + omega = 1.0 / (10000.0 ** (torch.arange(quarter, device=device, dtype=torch.float32) / quarter)) + yy, xx = torch.meshgrid( + torch.arange(h, device=device, dtype=torch.float32), + torch.arange(w, device=device, dtype=torch.float32), + indexing="ij", + ) + + def _emb(coord: torch.Tensor) -> torch.Tensor: # coord [h, w] -> [h*w, dim//2] + a = coord.reshape(-1, 1) * omega[None, :] + return torch.cat([a.sin(), a.cos()], dim=1) + + return torch.cat([_emb(yy), _emb(xx)], dim=1).to(dtype) # [h*w, dim] + + +class SimpleFPN(nn.Module): + """ViTDet-style multi-scale pyramid from a single ViT feature map. + + Input : native patch features ``[B, hidden, Hn, Wn]`` (stride-14 grid). + Output : list of ``[B, hidden, H_l, W_l]`` at each scale. ``scale > 1`` upsamples + (the fine level that buys small-object recall), ``== 1`` refines, ``< 1`` + downsamples. All levels stay hidden-dim so the cross-attention bank is cheap. + """ + + def __init__(self, config, scales: tuple[float, ...]): + super().__init__() + h = config.hidden_size + self.scales = tuple(scales) + self.blocks = nn.ModuleList() + for s in self.scales: + if s > 1: # upsample x2 -> fine (stride-7) + blk = nn.Sequential( + nn.ConvTranspose2d(h, h, kernel_size=2, stride=2), + LayerNorm2d(h), + _act(config.hidden_act), + ) + elif s == 1: # depthwise refine -> native + blk = nn.Sequential( + nn.Conv2d(h, h, kernel_size=3, padding=1, groups=h, bias=False), + LayerNorm2d(h), + ) + else: # downsample x2 -> coarse + blk = nn.Sequential( + nn.Conv2d(h, h, kernel_size=2, stride=2), + LayerNorm2d(h), + _act(config.hidden_act), + ) + self.blocks.append(blk) + + def forward(self, feat: torch.Tensor) -> list[torch.Tensor]: + return [blk(feat) for blk in self.blocks] + + +class MultiScaleBank(nn.Module): + """Flatten SimpleFPN levels into one token bank with level + 2D-sincos position. + + ``forward`` returns ``(feats, pos)`` each ``[B, sum(H_l*W_l), hidden]``: ``feats`` + are the level features (the cross-attn values), ``pos`` the additive positional + encoding (sincos per level + a learned per-level embedding) added to the keys. + """ + + def __init__(self, config, num_levels: int): + super().__init__() + self.level_embed = nn.Embedding(num_levels, config.hidden_size) + + def forward(self, levels: list[torch.Tensor]): + feats, poss = [], [] + for lvl, x in enumerate(levels): + b, c, h, w = x.shape + feats.append(x.flatten(2).transpose(1, 2)) # [B, h*w, c] + pos = _sincos_2d(h, w, c, x.device, x.dtype)[None] + self.level_embed.weight[lvl][None, None] + poss.append(pos.expand(b, -1, -1)) + return torch.cat(feats, dim=1), torch.cat(poss, dim=1) + + +class QueryCrossAttention(nn.Module): + """Additive cross-attention sublayer: queries attend to the multi-scale bank. + + Pre-norm + LayerScale + residual, mirroring :class:`Layer`. No FFN — the queries + already get self-attention + MLP from the shared stream they sit in; this only + injects multi-scale context. Cost is ``O(Q * N_bank)`` (Q ~= 100), so it stays + cheap even with a fine level in the bank — the point of B1. + """ + + def __init__(self, config): + super().__init__() + self.num_heads = config.num_attention_heads + self.head_dim = config.hidden_size // self.num_heads + self.scale = self.head_dim**-0.5 + self.norm_q = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.norm_kv = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.q_proj = nn.Linear(config.hidden_size, config.hidden_size) + self.k_proj = nn.Linear(config.hidden_size, config.hidden_size) + self.v_proj = nn.Linear(config.hidden_size, config.hidden_size) + self.out_proj = nn.Linear(config.hidden_size, config.hidden_size) + self.layer_scale = LayerScale(config) + + def forward(self, queries: torch.Tensor, bank_feats: torch.Tensor, bank_pos: torch.Tensor) -> torch.Tensor: + q_in = self.norm_q(queries) + kv = self.norm_kv(bank_feats) + b, q_n, _ = q_in.shape + n = kv.shape[1] + + def _split(x, length): + return x.view(b, length, self.num_heads, self.head_dim).transpose(1, 2) + + q = _split(self.q_proj(q_in), q_n) + k = _split(self.k_proj(kv + bank_pos), n) # position added to keys only (DETR) + v = _split(self.v_proj(kv), n) + attn = F.scaled_dot_product_attention(q, k, v, scale=self.scale) + attn = attn.transpose(1, 2).reshape(b, q_n, -1) + return queries + self.layer_scale(self.out_proj(attn)) + + class EoMTEncoder(nn.Module): """Pure-PyTorch ``EomtForUniversalSegmentation`` (forward + loss). @@ -280,6 +410,12 @@ class EoMTEncoder(nn.Module): ``family="detect"`` swaps the mask head for a :class:`BoxHead` and the mask/dice criterion for :class:`~eomt.box_loss.DetectionLoss`; masked attention is not used (there is no predicted mask to focus on), so the last blocks do full attention. + + When ``config.fpn_scales`` is set, a :class:`SimpleFPN` builds a multi-scale token + bank from the patch features at the query-injection point and each query block + gains a :class:`QueryCrossAttention` sublayer so queries can attend to it (the B1 + multi-scale fix). Left unset, the module is bit-for-bit the original EoMT — no + extra params, so it keeps state-dict parity with the HF model. """ def __init__(self, config, family: str = "instance"): @@ -316,6 +452,18 @@ def __init__(self, config, family: str = "instance"): self.criterion = EoMTLoss(config=config, weight_dict=self.weight_dict) self.register_buffer("attn_mask_probs", torch.ones(config.num_blocks)) + # Optional multi-scale (B1): SimpleFPN bank + one query→bank cross-attention + # per query block. Built only when enabled, so the default model adds no + # params and keeps state-dict parity with the HF EoMT. + fpn_scales = getattr(config, "fpn_scales", None) + self.fpn_scales = tuple(fpn_scales) if fpn_scales else None + if self.fpn_scales: + self.fpn = SimpleFPN(config, self.fpn_scales) + self.bank = MultiScaleBank(config, num_levels=len(self.fpn_scales)) + self.cross_blocks = nn.ModuleList( + [QueryCrossAttention(config) for _ in range(config.num_blocks)] + ) + # --- loss plumbing (mirrors HF get_loss_dict / get_loss) ---------------- def get_loss_dict(self, masks_queries_logits, class_queries_logits, mask_labels, class_labels): @@ -412,8 +560,17 @@ def forward( hidden_states = self.embeddings(pixel_values) + # Multi-scale bank (B1), built once from the patch features at the injection + # point and held fixed across the query blocks (Mask2Former-style memory). + bank_feats = bank_pos = None + for idx, layer_module in enumerate(self.layers): if idx == self.num_hidden_layers - self.config.num_blocks: + if self.fpn_scales: + num_prefix = self.embeddings.num_prefix_tokens + patch = hidden_states[:, num_prefix:] # [B, Hn*Wn, hidden] + patch2d = patch.transpose(1, 2).reshape(patch.shape[0], -1, *grid) + bank_feats, bank_pos = self.bank(self.fpn(patch2d)) query = ( self.query.weight[None, :, :] .expand(hidden_states.shape[0], -1, -1) @@ -442,6 +599,16 @@ def forward( hidden_states = layer_module(hidden_states, attention_mask) + # Multi-scale (B1): after the shared self-attention, let the query slice + # cross-attend to the bank so it gathers fine/coarse context the native + # grid alone can't provide. The next block's prediction reads the result. + if self.fpn_scales and in_query_blocks: + num_q = self.config.num_queries + refined = self.cross_blocks[block_idx]( + hidden_states[:, :num_q], bank_feats, bank_pos + ) + hidden_states = torch.cat((refined, hidden_states[:, num_q:]), dim=1) + sequence_output = self.layernorm(hidden_states) geom_final, class_queries_logits = self.predict(sequence_output, grid) geom_per_layer += (geom_final,) @@ -516,6 +683,7 @@ def build_model( aux_head_arch: dict | None = None, loss_weights: dict | None = None, num_upscale_blocks: int | None = None, + fpn_scales: tuple[float, ...] | list[float] | None = None, ) -> "EoMTModel": """Build an :class:`EoMTModel`. @@ -526,6 +694,8 @@ def build_model( criterion weights (mask/dice for instance, L1/GIoU for detect — see ``normalize_loss_weights``); ``num_upscale_blocks`` overrides the mask-head upsampling depth (instance only; ``None`` = size preset default). + ``fpn_scales`` enables the B1 multi-scale path (SimpleFPN + query + cross-attention); ``None`` = the original single-scale model. """ if family not in FAMILIES: raise ValueError(f"family={family!r} is not one of {FAMILIES}.") @@ -540,6 +710,7 @@ def build_model( aux_head_arch=aux_head_arch, loss_weights=loss_weights, num_upscale_blocks=num_upscale_blocks, + fpn_scales=fpn_scales, ) @@ -562,6 +733,7 @@ def __init__( aux_head_arch: dict | None = None, loss_weights: dict | None = None, num_upscale_blocks: int | None = None, + fpn_scales: tuple[float, ...] | list[float] | None = None, ): super().__init__() self.size = size @@ -599,6 +771,13 @@ def __init__( if family == "detect": config.l1_weight = float(self.loss_weights["l1_weight"]) config.giou_weight = float(self.loss_weights["giou_weight"]) + # Multi-scale (B1): stash the FPN scales on the config so ``EoMTEncoder`` + # (which only receives the config) builds the SimpleFPN + cross-attn modules. + # ``EomtConfig`` has no such field, so this is an added attribute; left unset + # the encoder is the original single-scale model. + self.fpn_scales = tuple(float(s) for s in fpn_scales) if fpn_scales else None + if self.fpn_scales: + config.fpn_scales = self.fpn_scales self.config = config # Effective upscale depth (resolved from the preset when not overridden). self.num_upscale_blocks = int(config.num_upscale_blocks) diff --git a/eomt/serialization.py b/eomt/serialization.py index e159063..dc4fde9 100644 --- a/eomt/serialization.py +++ b/eomt/serialization.py @@ -63,6 +63,7 @@ def wrap_checkpoint( letterbox: bool = False, loss_weights: dict | None = None, num_upscale_blocks: int | None = None, + fpn_scales: tuple[float, ...] | list[float] | None = None, norm_mean: Any = None, norm_std: Any = None, patch_size: int | None = None, @@ -104,6 +105,11 @@ def wrap_checkpoint( checkpoint["loss_weights"] = dict(loss_weights) if num_upscale_blocks is not None: checkpoint["num_upscale_blocks"] = int(num_upscale_blocks) + # Multi-scale (B1): load-bearing metadata — the conv weights don't encode which + # level is upsample vs downsample, so the scales must be recorded to rebuild the + # SimpleFPN + cross-attn modules with matching shapes on reload. + if fpn_scales is not None: + checkpoint["fpn_scales"] = [float(s) for s in fpn_scales] # Records the torchao recipe (if the weights are compressed) so ``load_model`` # re-creates the quantized layout before loading the subclassed tensors. if compression is not None: @@ -217,6 +223,9 @@ def load_model(path: str | Path, *, device: str = "auto") -> EoMTModel: # clean load — recover it from the state dict when the metadata predates it. loss_weights = ckpt.get("loss_weights") num_upscale_blocks = ckpt.get("num_upscale_blocks") or _infer_num_upscale_blocks(state) + # Multi-scale (B1) modules: rebuilt from the recorded scales (or recovered + # best-effort from the state dict when the metadata predates the field). + fpn_scales = ckpt.get("fpn_scales") or _infer_fpn_scales(state) model = build_model( size, @@ -228,6 +237,7 @@ def load_model(path: str | Path, *, device: str = "auto") -> EoMTModel: aux_head_arch=aux_head_arch, loss_weights=loss_weights, num_upscale_blocks=num_upscale_blocks, + fpn_scales=fpn_scales, ) # Compressed checkpoints carry int8 tensor subclasses (torchao). Re-create the # quantized layout, then load with ``assign=True`` so the saved subclassed @@ -327,6 +337,30 @@ def _infer_num_upscale_blocks(state: dict) -> int | None: return (max(idxs) + 1) if idxs else None +def _infer_fpn_scales(state: dict) -> list[float] | None: + """Best-effort recovery of B1 multi-scale: ``None`` unless FPN keys are present. + + The conv weights record only the *count* of pyramid levels (``eomt.fpn.blocks.``), + not their scale factors, so this falls back to the default scale tuple for that + count. Trust the ``fpn_scales`` metadata when present; this only covers a + metadata-stripped FPN checkpoint and warns implicitly via the level-count match. + """ + from .model import DEFAULT_FPN_SCALES + + idxs = set() + for key in state: + for prefix in ("eomt.fpn.blocks.", "fpn.blocks."): + if key.startswith(prefix): + head = key[len(prefix) :].split(".", 1)[0] + if head.isdigit(): + idxs.add(int(head)) + break + if not idxs: + return None + n = max(idxs) + 1 + return list(DEFAULT_FPN_SCALES) if n == len(DEFAULT_FPN_SCALES) else None + + def _infer_imgsz(state: dict, patch_size: int = 14) -> int: import math diff --git a/pyproject.toml b/pyproject.toml index 2836b37..8dd967f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "attr-eomt" -version = "0.1.3" +version = "1.0.0" description = "Standalone EoMT (Encoder-only Mask Transformer) for instance segmentation, with DINOv2 init, COCO training/validation and inference." readme = "README.md" requires-python = ">=3.10" diff --git a/scripts/train.py b/scripts/train.py index 6684af4..75a2164 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -33,8 +33,18 @@ def main() -> None: p.add_argument("--name", default=None, help="Run name (default: eomt-{size}).") p.add_argument("--weights", default=None, help="Init from a checkpoint/run to fine-tune (warm start).") p.add_argument("--resume", default=None, help="Resume a run from a checkpoint/run folder.") + p.add_argument( + "--fpn-scales", default="2,1,0.5", + help="B1 multi-scale SimpleFPN scales relative to the native grid (default " + "'2,1,0.5', on by default). Pass 'none'/'off' for the single-scale model.", + ) args = p.parse_args() + if args.fpn_scales.strip().lower() in ("", "none", "off", "0"): + fpn_scales = None + else: + fpn_scales = [float(s) for s in args.fpn_scales.split(",") if s.strip()] + # Init from a checkpoint (fine-tune / resume) or from a size (fresh, DINOv2 backbone). model = EoMT(args.weights or args.resume or args.size, device=args.device) result = model.train( @@ -45,6 +55,7 @@ def main() -> None: imgsz=args.imgsz, name=args.name, resume=bool(args.resume), + fpn_scales=fpn_scales, ) if result["best_metric"] >= 0: metric = "bbox mAP" if args.task == "detect" else "segm mAP"