File size: 4,624 Bytes
40d7c69 5878871 40d7c69 5878871 40d7c69 5878871 40d7c69 5878871 40d7c69 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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)
|