Spaces:
Running on Zero
Running on Zero
| """Shared DPT trunk. | |
| Same reassemble + RefineNet fusion as v1/v2, with one deliberate change: it | |
| returns features at **half** the input resolution and lets the heads run there. | |
| v2 upsampled to full resolution inside the trunk, so Head B materialised a | |
| (B, 128, 512, 512) bin-logit tensor — 3.4 GB at batch 12 in bf16, for no accuracy | |
| benefit. DPT's own design predicts at 1/2 and upsamples the *output*. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from torch.utils.checkpoint import checkpoint | |
| class ResidualConvUnit(nn.Module): | |
| def __init__(self, c: int): | |
| super().__init__() | |
| self.conv1 = nn.Conv2d(c, c, 3, padding=1, bias=False) | |
| self.bn1 = nn.GroupNorm(8, c) | |
| self.conv2 = nn.Conv2d(c, c, 3, padding=1, bias=False) | |
| self.bn2 = nn.GroupNorm(8, c) | |
| def forward(self, x): | |
| y = self.conv1(F.relu(x, inplace=False)) | |
| y = self.conv2(F.relu(self.bn1(y), inplace=True)) | |
| return x + self.bn2(y) | |
| class FeatureFusionBlock(nn.Module): | |
| """`use_skip=False` drops `rcu1` entirely. | |
| The deepest block is called with no skip connection, so its `rcu1` never | |
| takes part in the forward. On one GPU that is only six wasted parameters; | |
| under DDP with `find_unused_parameters=False` the reducer waits forever for | |
| gradients that never arrive and the next forward dies with "Expected to have | |
| finished reduction in the prior iteration" (indices 60-65 = `fuse[3].rcu1`). | |
| """ | |
| def __init__(self, c: int, use_skip: bool = True): | |
| super().__init__() | |
| self.rcu1 = ResidualConvUnit(c) if use_skip else None | |
| self.rcu2 = ResidualConvUnit(c) | |
| self.out = nn.Conv2d(c, c, 1) | |
| def forward(self, x, skip=None): | |
| if skip is not None: | |
| if self.rcu1 is None: | |
| raise RuntimeError("this FeatureFusionBlock was built without a " | |
| "skip path (use_skip=False)") | |
| s = self.rcu1(skip) | |
| # Align to the skip before adding. `resample[3]` is a stride-2 conv, | |
| # which ROUNDS UP (ceil), while the upsample below is a flat x2 — so | |
| # an odd token grid comes back one pixel too large and the add dies | |
| # with "The size of tensor a (6) must match the size of tensor b (5)". | |
| # A 512 px tile is a 32-grid and every TTA scale keeps it even, which | |
| # is why only small tile_size configs (a 64 px test tile at TTA 1.25 | |
| # -> an 80 px input -> a 5-grid) ever hit it. On an even grid the | |
| # shapes already match and this is a no-op, so nothing about the real | |
| # run changes. | |
| if x.shape[-2:] != s.shape[-2:]: | |
| x = F.interpolate(x, size=s.shape[-2:], mode="bilinear", | |
| align_corners=False) | |
| x = x + s | |
| x = self.rcu2(x) | |
| x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False) | |
| return self.out(x) | |
| class DPTTrunk(nn.Module): | |
| def __init__(self, in_ch: int, dim: int, grad_checkpoint: bool = False): | |
| super().__init__() | |
| self.dim = dim | |
| self.grad_checkpoint = grad_checkpoint | |
| proj = [96, 192, 384, 768] | |
| self.proj = nn.ModuleList(nn.Conv2d(in_ch, p, 1) for p in proj) | |
| self.resample = nn.ModuleList([ | |
| nn.ConvTranspose2d(proj[0], proj[0], 4, stride=4), # /16 -> /4 | |
| nn.ConvTranspose2d(proj[1], proj[1], 2, stride=2), # /16 -> /8 | |
| nn.Identity(), # /16 | |
| nn.Conv2d(proj[3], proj[3], 3, stride=2, padding=1), # /16 -> /32 | |
| ]) | |
| self.to_dim = nn.ModuleList(nn.Conv2d(p, dim, 3, padding=1, bias=False) for p in proj) | |
| # fuse[3] is the deepest block and runs without a skip — no rcu1. | |
| self.fuse = nn.ModuleList(FeatureFusionBlock(dim, use_skip=i < 3) | |
| for i in range(4)) | |
| def forward(self, feats: list[torch.Tensor]) -> torch.Tensor: | |
| f = [self.to_dim[i](self.resample[i](self.proj[i](feats[i]))) for i in range(4)] | |
| def run(fn, *a): | |
| if self.grad_checkpoint and self.training: | |
| return checkpoint(fn, *a, use_reentrant=False) | |
| return fn(*a) | |
| x = run(self.fuse[3], f[3]) # /32 -> /16 | |
| x = run(self.fuse[2], x, f[2]) # /16 -> /8 | |
| x = run(self.fuse[1], x, f[1]) # /8 -> /4 | |
| x = run(self.fuse[0], x, f[0]) # /4 -> /2 | |
| return x # (B, dim, H/2, W/2) | |