Spaces:

akashch1512
/

SingleViewHeigthEstimation

Sleeping

SingleViewHeigthEstimation / models /heads.py

fix: match v5 checkpoint architecture and benchmark scaled inference

5878871 verified about 1 month ago

19.6 kB

"""Heads + gated fusion + `DepthWizardNet`.
  Head A  metric nDSM regression   — softplus, not ReLU.  v2's `F.relu` on the raw
          conv output has exactly zero gradient wherever the head predicts a
          negative number, so a unit that drifts negative early stops learning.
  Head B  adaptive-bin classifier  — AdaBins-lite; owns the long tail.
  Head C  semantic segmentation    — auxiliary regulariser + the ground mask that
          `infer/calibrate.py` needs to tie an rDSM to a DEM.
  Fusion  per-pixel alpha over (A, B).
All heads run at H/2 and the final maps are upsampled once, at the end.
v4 adds one output: **`b_std`**, the standard deviation of Head B's bin
distribution, in metres.  It is a free per-pixel uncertainty — the bin head
already computes the distribution, so the second moment costs one extra
reduction — and the mean-teacher branch uses it to decide which of the teacher's
pseudo-labels on unlabeled Indian imagery are worth believing.  That also gives
Head B a second job, which is another reason for it not to collapse.
"""

from __future__ import annotations

from contextlib import nullcontext

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

from .dpt import DPTTrunk
from .encoder import DINOv3Encoder


def _conv_block(dim: int, out: int) -> nn.Sequential:
    return nn.Sequential(
        nn.Conv2d(dim, dim // 2, 3, padding=1, bias=False),
        nn.GroupNorm(8, dim // 2), nn.ReLU(True),
        nn.Conv2d(dim // 2, out, 1),
    )


class HeadA(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.net = _conv_block(dim, 1)
        # start near zero height: most pixels are ground
        nn.init.constant_(self.net[-1].bias, -2.0)

    def forward(self, x):
        return F.softplus(self.net(x))


class HeadB(nn.Module):
    """Global adaptive bin widths + per-pixel bin logits.
    `bin_centres` is the one path in this network with **no normalisation
    anywhere**: a global average pool straight into an unnormalised `nn.Linear`,
    whose output is then softmaxed.  That is what ended the second v4 Kaggle run,
    and the log says so precisely.  From e4 s1000 on, every line reads
        loss=nan (l1=nan sil=nan nrm=nan flat=nan bin=3.77 ent=2.97 seg=0.00 a=nan)
    and that combination has exactly one source.  `bin` and `ent` are computed
    from `b_logits`, which comes off `pixel_logits` — a `_conv_block`, so its
    first conv is followed by a GroupNorm and it stayed finite.  `bin` also reads
    `centres`, but only through `searchsorted` -> `clamp(0, K-1)`, which launders
    a NaN edge into a valid integer index.  Everything that went NaN — `l1`,
    `silog`, `normal`, `flat` and the reported `alpha` — reads `fused`, and
    `fused` reads `b["height"] = einsum(centres, probs)`.  So `centres` was NaN
    while `logits` was finite: the failure is in this method and nowhere else.
    Three things are fixed here, none of which change any parameter name, so
    existing checkpoints still load:
    1. **fp32.**  Under `torch.autocast` the `Linear` ran in fp16 on a pooled
       feature the trunk never bounds (`DPTTrunk` ends on a bare 1x1 conv).  One
       overflow to `inf` makes `softmax` return NaN for the whole row, and from
       there it is in `centres`, `fused`, and the gradient of every regression
       term.
    2. **RMS-normalised input.**  The global pool concentrates the gradient of
       every pixel in the tile onto this one small `Linear`, which is why it is
       the first thing to run away when the encoder unfreezes and the pooled
       feature scale shifts under it.  e3 is the log of that happening: `bin`
       leaves its 2.2-2.6 band at s1400 and thrashes between 2.6 and 6.74 —
       *above* `ln(96) = 4.56`, i.e. the head doing worse than a uniform
       prediction — for the rest of the epoch, because the bin edges themselves
       were moving faster than the pixel logits could track them.
    3. **A width floor.**  `torch.searchsorted` requires strictly increasing
       boundaries.  A saturated softmax gives bins whose width underflows to
       zero, so the edges are merely non-decreasing and the target index it
       returns is not defined.  Mixing in a little uniform mass keeps every bin
       open without meaningfully constraining the shape.
    The second moment is computed in fp32 for a separate reason: `std` is
    `sqrt(c2 - height**2)` and at the top of a 120 m range those two terms are
    both ~1.4e4, where fp16 has a resolution of 8.  The cancellation left `b_std`
    with ~3 m of pure quantisation noise, and `consistency_loss` gates the
    mean-teacher's pseudo-labels at `conf_m = 1.5` m — so the gate was being
    driven by rounding error rather than by the head's actual uncertainty.
    """

    # Softmax over 96 bins saturates long before this; the clamp only exists so
    # a diverging `Linear` cannot reach `inf` and take `softmax` with it.
    WIDTH_LOGIT_CLAMP = 15.0
    # Uniform mass mixed into the width distribution.  Each bin ends up at least
    # `WIDTH_FLOOR / K` of the range wide, so the edges are strictly increasing.
    WIDTH_FLOOR = 0.05

    def __init__(self, dim: int, n_bins: int, hmin: float, hmax: float):
        super().__init__()
        self.n_bins, self.hmin, self.hmax = n_bins, hmin, hmax
        self.readout = "mean"
        self.bin_widths = nn.Sequential(
            nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(dim, n_bins))
        self.pixel_logits = _conv_block(dim, n_bins)

    def bin_centres(self, x) -> torch.Tensor:
        # The pool is spelled out rather than run through `bin_widths[:2]` so it
        # can accumulate in fp32 *without* an `x.float()` first — that cast would
        # materialise a second copy of the whole (B, dim, H/2, W/2) trunk output,
        # ~134 MB per micro-batch at 512 px, and hold it live for the backward on
        # a profile whose own log reports 13-14 of 15 GiB in use.  `mean` over the
        # spatial dims with `dtype=` is exactly `AdaptiveAvgPool2d(1)` followed by
        # `Flatten`, at fp32 precision and fp16 footprint.  The two layers stay in
        # the Sequential anyway: they hold no parameters, so DDP does not see
        # them, and keeping them there is what keeps the trained `Linear` at
        # `bin_widths.2.*` so existing checkpoints load unchanged.
        with torch.autocast(x.device.type, enabled=False):
            g = x.mean(dim=(-2, -1), dtype=torch.float32)           # (B, C)
            g = g * torch.rsqrt(g.pow(2).mean(dim=1, keepdim=True) + 1e-6)
            logits = self.bin_widths[2](g).clamp(-self.WIDTH_LOGIT_CLAMP,
                                                 self.WIDTH_LOGIT_CLAMP)
            p = F.softmax(logits, dim=1)
            p = (p + self.WIDTH_FLOOR / self.n_bins) / (1.0 + self.WIDTH_FLOOR)
            w = p * (self.hmax - self.hmin)
            edges = F.pad(self.hmin + torch.cumsum(w, dim=1), (1, 0), value=self.hmin)
            return 0.5 * (edges[:, 1:] + edges[:, :-1])             # (B, K)

    def forward(self, x):
        centres = self.bin_centres(x)                               # fp32
        logits = self.pixel_logits(x)                               # (B, K, h, w)
        # Both moments in fp32: `einsum` is on autocast's fp16 list, and `c2`
        # reaches `hmax ** 2` = 1.4e4 where fp16 cannot resolve the difference
        # that `std` is made of.
        with torch.autocast(x.device.type, enabled=False):
            probs = F.softmax(logits.float(), dim=1)
            height = torch.einsum("bk,bkhw->bhw", centres, probs).unsqueeze(1)
            c2 = torch.einsum("bk,bkhw->bhw", centres ** 2, probs).unsqueeze(1)
            std = (c2 - height ** 2).clamp_min(0.0).sqrt()
            if self.readout == "single_mode":
                # Keep the descending basin around the most probable bin.
                # Stop at the first rise toward another mode on either side.
                peak = probs.argmax(1, keepdim=True)
                idx = torch.arange(self.n_bins, device=x.device)[None, :, None, None]
                next_p = torch.cat([probs[:, 1:], probs[:, -1:]], 1)
                prev_p = torch.cat([probs[:, :1], probs[:, :-1]], 1)
                left_break = (idx < peak) & (probs > next_p)
                right_break = (idx > peak) & (probs > prev_p)
                keep_left = left_break.flip(1).cumsum(1).flip(1) == 0
                keep_right = right_break.cumsum(1) == 0
                local = probs * (keep_left & keep_right)
                local = local / local.sum(1, keepdim=True).clamp_min(1e-12)
                height = torch.einsum("bk,bkhw->bhw", centres, local).unsqueeze(1)
        return {"height": height, "logits": logits, "centres": centres, "std": std}


class GatedFusion(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.gate = nn.Sequential(
            nn.Conv2d(dim + 2, dim // 4, 3, padding=1, bias=False),
            nn.GroupNorm(8, dim // 4), nn.ReLU(True),
            nn.Conv2d(dim // 4, 1, 1),
        )

    def forward(self, feat, a, b, scale=None):
        # `b` arrives in fp32 (HeadB computes its moments there — see HeadB's
        # docstring) while `feat` and `a` are autocast fp16.  `torch.cat` type-
        # promotes, so concatenating them as-is would silently widen all `dim`
        # channels of `feat` to fp32 — ~70 MB per micro-batch at 512 px, on a
        # profile the log shows sitting at 13-14 of 15 GiB.  The gate is a conv
        # stack autocast would run in fp16 regardless, so the cast is free to
        # make explicit and narrow.
        #
        # `scale` is the per-sample factor `feat` was already multiplied by (see
        # `DepthWizardNet.forward`).  The gate's first conv has no bias and feeds
        # a GroupNorm, so scaling its whole input leaves alpha unchanged, but
        # only if `a` and `b` go in scaled along with `feat`.
        ga, gb = (a, b) if scale is None else (a * scale, b * scale)
        g = torch.cat([feat, ga.to(feat.dtype), gb.to(feat.dtype)], dim=1)
        alpha = torch.sigmoid(self.gate(g))
        # The blend itself stays in the wider of the two: this is the metric
        # height that every regression loss is computed against.
        alpha = alpha.to(b.dtype)
        return alpha * a.to(b.dtype) + (1.0 - alpha) * b, alpha


# ---------------------------------------------------------------------
# v5: the full-resolution detail path
# ---------------------------------------------------------------------
class DetailStem(nn.Module):
    """A shallow conv stem on the RGB itself, at full and half resolution.
    The ViT sees 16 px patches and the DPT trunk stops at H/2, so nothing in v4
    ever looked at the image above 1/4 resolution with more than one conv —
    roof edges, tree crowns and walls were reconstructed from patch tokens and
    then bilinearly upsampled.  This gives the heads (at H/2) and the upsampler
    (to H) direct access to image edges.
    """

    def __init__(self, c1: int = 32, c2: int = 64):
        super().__init__()
        # not `full` / `half`: `nn.Module.half` is the fp16 cast
        self.s1 = nn.Sequential(nn.Conv2d(3, c1, 3, padding=1, bias=False),
                                nn.GroupNorm(8, c1), nn.ReLU(True))
        self.s2 = nn.Sequential(nn.Conv2d(c1, c2, 3, stride=2, padding=1, bias=False),
                                nn.GroupNorm(8, c2), nn.ReLU(True))

    def forward(self, image):
        f = self.s1(image)
        return f, self.s2(f)


def _bilinear_up2_logits() -> torch.Tensor:
    """(36,) logits that make `ConvexUp2x` reproduce bilinear 2x upsampling.
    Channel c = k*4 + di*2 + dj, neighbour k = dy*3 + dx at offset (dy-1, dx-1),
    output sub-pixel (di, dj).  With align_corners=False, sub-pixel 0 sits a
    quarter pixel before the source centre: weights 0.25 / 0.75 on rows i-1 / i;
    sub-pixel 1 uses rows i / i+1 at 0.75 / 0.25.
    """
    w1 = {0: {-1: 0.25, 0: 0.75}, 1: {0: 0.75, 1: 0.25}}
    out = torch.full((9, 2, 2), 1e-4)
    for dy in range(3):
        for dx in range(3):
            for di in range(2):
                for dj in range(2):
                    wy = w1[di].get(dy - 1, 0.0)
                    wx = w1[dj].get(dx - 1, 0.0)
                    if wy * wx > 0:
                        out[dy * 3 + dx, di, dj] = wy * wx
    return out.log().reshape(36)


class ConvexUp2x(nn.Module):
    """RAFT-style learned upsampling, H/2 -> H: every output pixel is a convex
    combination of its 3x3 low-res neighbourhood, weights predicted from the
    trunk features and the full-resolution stem.  Where the image has an edge
    the weights can snap to one side of it instead of averaging across — which
    is what a roof edge or a wall needs and what bilinear cannot do.
    The last conv is zero-weighted with the bilinear logits as its bias, so at
    initialisation this *is* bilinear upsampling (to ~1e-4) and a v4 checkpoint
    loaded into a v5 net predicts what it predicted before.  The neighbourhood
    is gathered with pad + 9 slices rather than `F.unfold`, which keeps the ONNX
    graph to Pad / Slice / Softmax / DepthToSpace-free reshapes.
    """

    def __init__(self, cin: int, hidden: int = 128):
        super().__init__()
        self.net = nn.Sequential(nn.Conv2d(cin, hidden, 3, padding=1), nn.ReLU(True),
                                 nn.Conv2d(hidden, 36, 1))
        nn.init.zeros_(self.net[-1].weight)
        with torch.no_grad():
            self.net[-1].bias.copy_(_bilinear_up2_logits())

    def weights(self, feat):
        b, _, h, w = feat.shape
        return self.net(feat).float().view(b, 1, 9, 2, 2, h, w).softmax(2)

    @staticmethod
    def apply(t, m):
        """t (B, C, h, w), m from `weights` -> (B, C, 2h, 2w)."""
        b, c, h, w = t.shape
        tp = F.pad(t.float(), (1, 1, 1, 1), mode="replicate")
        nb = torch.stack([tp[..., dy:dy + h, dx:dx + w]
                          for dy in range(3) for dx in range(3)], 2)   # (B, C, 9, h, w)
        out = (m * nb.view(b, c, 9, 1, 1, h, w)).sum(2)                 # (B, C, 2, 2, h, w)
        return out.permute(0, 1, 4, 2, 5, 3).reshape(b, c, 2 * h, 2 * w)


class DepthWizardNet(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        from config import N_SEG_CLASSES

        self.cfg = cfg
        self.encoder = DINOv3Encoder(cfg)
        d = cfg.decoder_dim
        self.trunk = DPTTrunk(self.encoder.hidden, d, cfg.grad_checkpoint_decoder)
        self.head_a = HeadA(d)
        self.head_b = HeadB(d, cfg.n_bins, cfg.bin_min_m, cfg.bin_max_m)
        self.head_c = _conv_block(d, N_SEG_CLASSES)
        self.fusion = GatedFusion(d)
        self.detail = bool(getattr(cfg, "detail_branch", False))
        if self.detail:
            c2 = int(getattr(cfg, "detail_dim", 64))
            self.stem = DetailStem(max(8, c2 // 2), c2)
            # zero-init: at t=0 the trunk features pass through unchanged
            self.inject = nn.Conv2d(d + c2, d, 1)
            nn.init.zeros_(self.inject.weight)
            nn.init.zeros_(self.inject.bias)
            self.upsampler = ConvexUp2x(d + c2 + 4 * max(8, c2 // 2))

    # Largest |feature| handed to the fp16 heads.  Their first conv sums
    # 9 x 256 products, so this keeps its output far inside fp16's 65504.
    FP16_FEAT_MAX = 64.0

    def forward(self, image: torch.Tensor) -> dict:
        """Under fp16 autocast, the parts with no normalisation run in fp32.
        bf16, fp32 and ONNX export take the same path as before.
        The DPT trunk is a residual stream with no normalisation, and neither
        the detail `inject` nor the upsampler's first conv (bias, then ReLU) has
        any either.  v4 was trained in bf16 on the H100, and its stream grows
        past fp16's 65504 on some tiles.  The first T4 warm start shows what
        follows: `inf` in the trunk output (NON-FINITE b_centres), NaN on
        dfc23 val, and 54 % then 97 % of optimiser steps skipped, because one
        bad micro-batch spoils all 16 in its step.
        The heads, the fusion gate and HeadB's bin-width pool each start with
        either a bias-free conv followed by GroupNorm or an RMS normalisation.
        That makes them invariant to a positive per-sample scale on their input.
        So they stay in fp16 and get a copy of the features scaled down to
        FP16_FEAT_MAX.  The result is exact apart from GroupNorm's eps.
        """
        hw = image.shape[-2:]
        dev = image.device.type
        fp16 = (torch.is_autocast_enabled(dev)
                and torch.get_autocast_dtype(dev) == torch.float16)
        fp32 = (lambda: torch.autocast(dev, enabled=False)) if fp16 else nullcontext
        feats = self.encoder(image)
        with fp32():
            x = self.trunk([f.float() for f in feats] if fp16 else feats)  # (B, d, H/2, W/2)
        up = lambda t: F.interpolate(t, size=hw, mode="bilinear", align_corners=False)  # noqa: E731
        up_h = up
        if self.detail:
            s_full, s_half = self.stem(image)
            if s_half.shape[-2:] != x.shape[-2:]:
                s_half = F.interpolate(s_half, size=x.shape[-2:], mode="bilinear",
                                       align_corners=False)
            with fp32():
                x = x + self.inject(torch.cat([x, s_half.to(x.dtype)], 1))
                if tuple(hw) == (2 * x.shape[-2], 2 * x.shape[-1]):
                    m = self.upsampler.weights(torch.cat(
                        [x, s_half.to(x.dtype), F.pixel_unshuffle(s_full, 2).to(x.dtype)], 1))
                    up_h = lambda t: ConvexUp2x.apply(t, m)  # noqa: E731
        scale = None
        if fp16:
            # Detached, because the heads are blind to it: no gradient flows through it.
            peak = x.detach().abs().amax(dim=(1, 2, 3), keepdim=True)
            scale = (self.FP16_FEAT_MAX / peak.clamp_min(1e-12)).clamp_max(1.0)
            x = (x * scale).to(torch.float16)
        a = self.head_a(x)
        b = self.head_b(x)
        fused, alpha = self.fusion(x, a, b["height"], scale)
        if getattr(self, "eval_head_b_weight", None) is not None:
            weight = self.eval_head_b_weight
            fused = (1.0 - weight) * a.float() + weight * b["height"]

        return {
            "a": up_h(a),
            "b": up_h(b["height"]),
            "fused": up_h(fused),
            "alpha": up(alpha),
            "b_std": up(b["std"]),
            "seg": up(self.head_c(x)),
            # kept at half res — the bin loss is computed against a downsampled
            # target rather than blowing a (B, K, H, W) tensor up to full size.
            "b_logits": b["logits"],
            "b_centres": b["centres"],
        }

    def decoder_parameters(self):
        for n, p in self.named_parameters():
            if p.requires_grad and not n.startswith("encoder."):
                yield p

    def head_state_dict(self) -> dict:
        """Everything except the public frozen encoder backbone.
        Encoder weights are only included once they have actually been trained,
        which `train.py` signals with `include_encoder=True`.
        """
        return {k: v for k, v in self.state_dict().items()
                if not k.startswith("encoder.model.")}

    def full_state_dict(self) -> dict:
        return self.state_dict()


# v3 checkpoints and any external code that imported the old name keep working.
DepthWizardNetV3 = DepthWizardNet

	"""Heads + gated fusion + `DepthWizardNet`.

	Head A metric nDSM regression — softplus, not ReLU. v2's `F.relu` on the raw
	conv output has exactly zero gradient wherever the head predicts a
	negative number, so a unit that drifts negative early stops learning.
	Head B adaptive-bin classifier — AdaBins-lite; owns the long tail.
	Head C semantic segmentation — auxiliary regulariser + the ground mask that
	`infer/calibrate.py` needs to tie an rDSM to a DEM.
	Fusion per-pixel alpha over (A, B).

	All heads run at H/2 and the final maps are upsampled once, at the end.

	v4 adds one output: `b_std`, the standard deviation of Head B's bin
	distribution, in metres. It is a free per-pixel uncertainty — the bin head
	already computes the distribution, so the second moment costs one extra
	reduction — and the mean-teacher branch uses it to decide which of the teacher's
	pseudo-labels on unlabeled Indian imagery are worth believing. That also gives
	Head B a second job, which is another reason for it not to collapse.
	"""

	from __future__ import annotations

	from contextlib import nullcontext

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

	from .dpt import DPTTrunk
	from .encoder import DINOv3Encoder


	def _conv_block(dim: int, out: int) -> nn.Sequential:
	return nn.Sequential(
	nn.Conv2d(dim, dim // 2, 3, padding=1, bias=False),
	nn.GroupNorm(8, dim // 2), nn.ReLU(True),
	nn.Conv2d(dim // 2, out, 1),
	)


	class HeadA(nn.Module):
	def __init__(self, dim: int):
	super().__init__()
	self.net = _conv_block(dim, 1)
	# start near zero height: most pixels are ground
	nn.init.constant_(self.net[-1].bias, -2.0)

	def forward(self, x):
	return F.softplus(self.net(x))


	class HeadB(nn.Module):
	"""Global adaptive bin widths + per-pixel bin logits.

	`bin_centres` is the one path in this network with **no normalisation
	anywhere**: a global average pool straight into an unnormalised `nn.Linear`,
	whose output is then softmaxed. That is what ended the second v4 Kaggle run,
	and the log says so precisely. From e4 s1000 on, every line reads

	loss=nan (l1=nan sil=nan nrm=nan flat=nan bin=3.77 ent=2.97 seg=0.00 a=nan)

	and that combination has exactly one source. `bin` and `ent` are computed
	from `b_logits`, which comes off `pixel_logits` — a `_conv_block`, so its
	first conv is followed by a GroupNorm and it stayed finite. `bin` also reads
	`centres`, but only through `searchsorted` -> `clamp(0, K-1)`, which launders
	a NaN edge into a valid integer index. Everything that went NaN — `l1`,
	`silog`, `normal`, `flat` and the reported `alpha` — reads `fused`, and
	`fused` reads `b["height"] = einsum(centres, probs)`. So `centres` was NaN
	while `logits` was finite: the failure is in this method and nowhere else.

	Three things are fixed here, none of which change any parameter name, so
	existing checkpoints still load:

	1. fp32. Under `torch.autocast` the `Linear` ran in fp16 on a pooled
	feature the trunk never bounds (`DPTTrunk` ends on a bare 1x1 conv). One
	overflow to `inf` makes `softmax` return NaN for the whole row, and from
	there it is in `centres`, `fused`, and the gradient of every regression
	term.
	2. RMS-normalised input. The global pool concentrates the gradient of
	every pixel in the tile onto this one small `Linear`, which is why it is
	the first thing to run away when the encoder unfreezes and the pooled
	feature scale shifts under it. e3 is the log of that happening: `bin`
	leaves its 2.2-2.6 band at s1400 and thrashes between 2.6 and 6.74 —
	above `ln(96) = 4.56`, i.e. the head doing worse than a uniform
	prediction — for the rest of the epoch, because the bin edges themselves
	were moving faster than the pixel logits could track them.
	3. A width floor. `torch.searchsorted` requires strictly increasing
	boundaries. A saturated softmax gives bins whose width underflows to
	zero, so the edges are merely non-decreasing and the target index it
	returns is not defined. Mixing in a little uniform mass keeps every bin
	open without meaningfully constraining the shape.

	The second moment is computed in fp32 for a separate reason: `std` is
	`sqrt(c2 - height**2)` and at the top of a 120 m range those two terms are
	both ~1.4e4, where fp16 has a resolution of 8. The cancellation left `b_std`
	with ~3 m of pure quantisation noise, and `consistency_loss` gates the
	mean-teacher's pseudo-labels at `conf_m = 1.5` m — so the gate was being
	driven by rounding error rather than by the head's actual uncertainty.
	"""

	# Softmax over 96 bins saturates long before this; the clamp only exists so
	# a diverging `Linear` cannot reach `inf` and take `softmax` with it.
	WIDTH_LOGIT_CLAMP = 15.0
	# Uniform mass mixed into the width distribution. Each bin ends up at least
	# `WIDTH_FLOOR / K` of the range wide, so the edges are strictly increasing.
	WIDTH_FLOOR = 0.05

	def __init__(self, dim: int, n_bins: int, hmin: float, hmax: float):
	super().__init__()
	self.n_bins, self.hmin, self.hmax = n_bins, hmin, hmax
	self.readout = "mean"
	self.bin_widths = nn.Sequential(
	nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(dim, n_bins))
	self.pixel_logits = _conv_block(dim, n_bins)

	def bin_centres(self, x) -> torch.Tensor:
	# The pool is spelled out rather than run through `bin_widths[:2]` so it
	# can accumulate in fp32 without an `x.float()` first — that cast would
	# materialise a second copy of the whole (B, dim, H/2, W/2) trunk output,
	# ~134 MB per micro-batch at 512 px, and hold it live for the backward on
	# a profile whose own log reports 13-14 of 15 GiB in use. `mean` over the
	# spatial dims with `dtype=` is exactly `AdaptiveAvgPool2d(1)` followed by
	# `Flatten`, at fp32 precision and fp16 footprint. The two layers stay in
	# the Sequential anyway: they hold no parameters, so DDP does not see
	# them, and keeping them there is what keeps the trained `Linear` at
	# `bin_widths.2.*` so existing checkpoints load unchanged.
	with torch.autocast(x.device.type, enabled=False):
	g = x.mean(dim=(-2, -1), dtype=torch.float32) # (B, C)
	g = g * torch.rsqrt(g.pow(2).mean(dim=1, keepdim=True) + 1e-6)
	logits = self.bin_widths[2](g).clamp(-self.WIDTH_LOGIT_CLAMP,
	self.WIDTH_LOGIT_CLAMP)
	p = F.softmax(logits, dim=1)
	p = (p + self.WIDTH_FLOOR / self.n_bins) / (1.0 + self.WIDTH_FLOOR)
	w = p * (self.hmax - self.hmin)
	edges = F.pad(self.hmin + torch.cumsum(w, dim=1), (1, 0), value=self.hmin)
	return 0.5 * (edges[:, 1:] + edges[:, :-1]) # (B, K)

	def forward(self, x):
	centres = self.bin_centres(x) # fp32
	logits = self.pixel_logits(x) # (B, K, h, w)
	# Both moments in fp32: `einsum` is on autocast's fp16 list, and `c2`
	# reaches `hmax ** 2` = 1.4e4 where fp16 cannot resolve the difference
	# that `std` is made of.
	with torch.autocast(x.device.type, enabled=False):
	probs = F.softmax(logits.float(), dim=1)
	height = torch.einsum("bk,bkhw->bhw", centres, probs).unsqueeze(1)
	c2 = torch.einsum("bk,bkhw->bhw", centres ** 2, probs).unsqueeze(1)
	std = (c2 - height ** 2).clamp_min(0.0).sqrt()
	if self.readout == "single_mode":
	# Keep the descending basin around the most probable bin.
	# Stop at the first rise toward another mode on either side.
	peak = probs.argmax(1, keepdim=True)
	idx = torch.arange(self.n_bins, device=x.device)[None, :, None, None]
	next_p = torch.cat([probs[:, 1:], probs[:, -1:]], 1)
	prev_p = torch.cat([probs[:, :1], probs[:, :-1]], 1)
	left_break = (idx < peak) & (probs > next_p)
	right_break = (idx > peak) & (probs > prev_p)
	keep_left = left_break.flip(1).cumsum(1).flip(1) == 0
	keep_right = right_break.cumsum(1) == 0
	local = probs * (keep_left & keep_right)
	local = local / local.sum(1, keepdim=True).clamp_min(1e-12)
	height = torch.einsum("bk,bkhw->bhw", centres, local).unsqueeze(1)
	return {"height": height, "logits": logits, "centres": centres, "std": std}


	class GatedFusion(nn.Module):
	def __init__(self, dim: int):
	super().__init__()
	self.gate = nn.Sequential(
	nn.Conv2d(dim + 2, dim // 4, 3, padding=1, bias=False),
	nn.GroupNorm(8, dim // 4), nn.ReLU(True),
	nn.Conv2d(dim // 4, 1, 1),
	)

	def forward(self, feat, a, b, scale=None):
	# `b` arrives in fp32 (HeadB computes its moments there — see HeadB's
	# docstring) while `feat` and `a` are autocast fp16. `torch.cat` type-
	# promotes, so concatenating them as-is would silently widen all `dim`
	# channels of `feat` to fp32 — ~70 MB per micro-batch at 512 px, on a
	# profile the log shows sitting at 13-14 of 15 GiB. The gate is a conv
	# stack autocast would run in fp16 regardless, so the cast is free to
	# make explicit and narrow.
	#
	# `scale` is the per-sample factor `feat` was already multiplied by (see
	# `DepthWizardNet.forward`). The gate's first conv has no bias and feeds
	# a GroupNorm, so scaling its whole input leaves alpha unchanged, but
	# only if `a` and `b` go in scaled along with `feat`.
	ga, gb = (a, b) if scale is None else (a * scale, b * scale)
	g = torch.cat([feat, ga.to(feat.dtype), gb.to(feat.dtype)], dim=1)
	alpha = torch.sigmoid(self.gate(g))
	# The blend itself stays in the wider of the two: this is the metric
	# height that every regression loss is computed against.
	alpha = alpha.to(b.dtype)
	return alpha * a.to(b.dtype) + (1.0 - alpha) * b, alpha


	# ---------------------------------------------------------------------
	# v5: the full-resolution detail path
	# ---------------------------------------------------------------------
	class DetailStem(nn.Module):
	"""A shallow conv stem on the RGB itself, at full and half resolution.

	The ViT sees 16 px patches and the DPT trunk stops at H/2, so nothing in v4
	ever looked at the image above 1/4 resolution with more than one conv —
	roof edges, tree crowns and walls were reconstructed from patch tokens and
	then bilinearly upsampled. This gives the heads (at H/2) and the upsampler
	(to H) direct access to image edges.
	"""

	def __init__(self, c1: int = 32, c2: int = 64):
	super().__init__()
	# not `full` / `half`: `nn.Module.half` is the fp16 cast
	self.s1 = nn.Sequential(nn.Conv2d(3, c1, 3, padding=1, bias=False),
	nn.GroupNorm(8, c1), nn.ReLU(True))
	self.s2 = nn.Sequential(nn.Conv2d(c1, c2, 3, stride=2, padding=1, bias=False),
	nn.GroupNorm(8, c2), nn.ReLU(True))

	def forward(self, image):
	f = self.s1(image)
	return f, self.s2(f)


	def _bilinear_up2_logits() -> torch.Tensor:
	"""(36,) logits that make `ConvexUp2x` reproduce bilinear 2x upsampling.

	Channel c = k4 + di2 + dj, neighbour k = dy*3 + dx at offset (dy-1, dx-1),
	output sub-pixel (di, dj). With align_corners=False, sub-pixel 0 sits a
	quarter pixel before the source centre: weights 0.25 / 0.75 on rows i-1 / i;
	sub-pixel 1 uses rows i / i+1 at 0.75 / 0.25.
	"""
	w1 = {0: {-1: 0.25, 0: 0.75}, 1: {0: 0.75, 1: 0.25}}
	out = torch.full((9, 2, 2), 1e-4)
	for dy in range(3):
	for dx in range(3):
	for di in range(2):
	for dj in range(2):
	wy = w1[di].get(dy - 1, 0.0)
	wx = w1[dj].get(dx - 1, 0.0)
	if wy * wx > 0:
	out[dy * 3 + dx, di, dj] = wy * wx
	return out.log().reshape(36)


	class ConvexUp2x(nn.Module):
	"""RAFT-style learned upsampling, H/2 -> H: every output pixel is a convex
	combination of its 3x3 low-res neighbourhood, weights predicted from the
	trunk features and the full-resolution stem. Where the image has an edge
	the weights can snap to one side of it instead of averaging across — which
	is what a roof edge or a wall needs and what bilinear cannot do.

	The last conv is zero-weighted with the bilinear logits as its bias, so at
	initialisation this is bilinear upsampling (to ~1e-4) and a v4 checkpoint
	loaded into a v5 net predicts what it predicted before. The neighbourhood
	is gathered with pad + 9 slices rather than `F.unfold`, which keeps the ONNX
	graph to Pad / Slice / Softmax / DepthToSpace-free reshapes.
	"""

	def __init__(self, cin: int, hidden: int = 128):
	super().__init__()
	self.net = nn.Sequential(nn.Conv2d(cin, hidden, 3, padding=1), nn.ReLU(True),
	nn.Conv2d(hidden, 36, 1))
	nn.init.zeros_(self.net[-1].weight)
	with torch.no_grad():
	self.net[-1].bias.copy_(_bilinear_up2_logits())

	def weights(self, feat):
	b, _, h, w = feat.shape
	return self.net(feat).float().view(b, 1, 9, 2, 2, h, w).softmax(2)

	@staticmethod
	def apply(t, m):
	"""t (B, C, h, w), m from `weights` -> (B, C, 2h, 2w)."""
	b, c, h, w = t.shape
	tp = F.pad(t.float(), (1, 1, 1, 1), mode="replicate")
	nb = torch.stack([tp[..., dy:dy + h, dx:dx + w]
	for dy in range(3) for dx in range(3)], 2) # (B, C, 9, h, w)
	out = (m * nb.view(b, c, 9, 1, 1, h, w)).sum(2) # (B, C, 2, 2, h, w)
	return out.permute(0, 1, 4, 2, 5, 3).reshape(b, c, 2 * h, 2 * w)


	class DepthWizardNet(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	from config import N_SEG_CLASSES

	self.cfg = cfg
	self.encoder = DINOv3Encoder(cfg)
	d = cfg.decoder_dim
	self.trunk = DPTTrunk(self.encoder.hidden, d, cfg.grad_checkpoint_decoder)
	self.head_a = HeadA(d)
	self.head_b = HeadB(d, cfg.n_bins, cfg.bin_min_m, cfg.bin_max_m)
	self.head_c = _conv_block(d, N_SEG_CLASSES)
	self.fusion = GatedFusion(d)
	self.detail = bool(getattr(cfg, "detail_branch", False))
	if self.detail:
	c2 = int(getattr(cfg, "detail_dim", 64))
	self.stem = DetailStem(max(8, c2 // 2), c2)
	# zero-init: at t=0 the trunk features pass through unchanged
	self.inject = nn.Conv2d(d + c2, d, 1)
	nn.init.zeros_(self.inject.weight)
	nn.init.zeros_(self.inject.bias)
	self.upsampler = ConvexUp2x(d + c2 + 4 * max(8, c2 // 2))

	# Largest \|feature\| handed to the fp16 heads. Their first conv sums
	# 9 x 256 products, so this keeps its output far inside fp16's 65504.
	FP16_FEAT_MAX = 64.0

	def forward(self, image: torch.Tensor) -> dict:
	"""Under fp16 autocast, the parts with no normalisation run in fp32.
	bf16, fp32 and ONNX export take the same path as before.

	The DPT trunk is a residual stream with no normalisation, and neither
	the detail `inject` nor the upsampler's first conv (bias, then ReLU) has
	any either. v4 was trained in bf16 on the H100, and its stream grows
	past fp16's 65504 on some tiles. The first T4 warm start shows what
	follows: `inf` in the trunk output (NON-FINITE b_centres), NaN on
	dfc23 val, and 54 % then 97 % of optimiser steps skipped, because one
	bad micro-batch spoils all 16 in its step.

	The heads, the fusion gate and HeadB's bin-width pool each start with
	either a bias-free conv followed by GroupNorm or an RMS normalisation.
	That makes them invariant to a positive per-sample scale on their input.
	So they stay in fp16 and get a copy of the features scaled down to
	FP16_FEAT_MAX. The result is exact apart from GroupNorm's eps.
	"""
	hw = image.shape[-2:]
	dev = image.device.type
	fp16 = (torch.is_autocast_enabled(dev)
	and torch.get_autocast_dtype(dev) == torch.float16)
	fp32 = (lambda: torch.autocast(dev, enabled=False)) if fp16 else nullcontext
	feats = self.encoder(image)
	with fp32():
	x = self.trunk([f.float() for f in feats] if fp16 else feats) # (B, d, H/2, W/2)
	up = lambda t: F.interpolate(t, size=hw, mode="bilinear", align_corners=False) # noqa: E731
	up_h = up
	if self.detail:
	s_full, s_half = self.stem(image)
	if s_half.shape[-2:] != x.shape[-2:]:
	s_half = F.interpolate(s_half, size=x.shape[-2:], mode="bilinear",
	align_corners=False)
	with fp32():
	x = x + self.inject(torch.cat([x, s_half.to(x.dtype)], 1))
	if tuple(hw) == (2 * x.shape[-2], 2 * x.shape[-1]):
	m = self.upsampler.weights(torch.cat(
	[x, s_half.to(x.dtype), F.pixel_unshuffle(s_full, 2).to(x.dtype)], 1))
	up_h = lambda t: ConvexUp2x.apply(t, m) # noqa: E731
	scale = None
	if fp16:
	# Detached, because the heads are blind to it: no gradient flows through it.
	peak = x.detach().abs().amax(dim=(1, 2, 3), keepdim=True)
	scale = (self.FP16_FEAT_MAX / peak.clamp_min(1e-12)).clamp_max(1.0)
	x = (x * scale).to(torch.float16)
	a = self.head_a(x)
	b = self.head_b(x)
	fused, alpha = self.fusion(x, a, b["height"], scale)
	if getattr(self, "eval_head_b_weight", None) is not None:
	weight = self.eval_head_b_weight
	fused = (1.0 - weight) * a.float() + weight * b["height"]

	return {
	"a": up_h(a),
	"b": up_h(b["height"]),
	"fused": up_h(fused),
	"alpha": up(alpha),
	"b_std": up(b["std"]),
	"seg": up(self.head_c(x)),
	# kept at half res — the bin loss is computed against a downsampled
	# target rather than blowing a (B, K, H, W) tensor up to full size.
	"b_logits": b["logits"],
	"b_centres": b["centres"],
	}

	def decoder_parameters(self):
	for n, p in self.named_parameters():
	if p.requires_grad and not n.startswith("encoder."):
	yield p

	def head_state_dict(self) -> dict:
	"""Everything except the public frozen encoder backbone.

	Encoder weights are only included once they have actually been trained,
	which `train.py` signals with `include_encoder=True`.
	"""
	return {k: v for k, v in self.state_dict().items()
	if not k.startswith("encoder.model.")}

	def full_state_dict(self) -> dict:
	return self.state_dict()


	# v3 checkpoints and any external code that imported the old name keep working.
	DepthWizardNetV3 = DepthWizardNet