Pikaliov
1b3206b6a7
- Add enhanced README with project structure and quick start guide - Initialize repository with DCNv4 CUDA extension (PyTorch module) - Include classification, detection, and segmentation subdirectories - Reference upstream OpenGVLab DCNv4 implementation Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
239 lines
9.7 KiB
Python
239 lines
9.7 KiB
Python
"""Strip-DCN feasibility gate (Phase A3, PLAN_code_stripDCN_stripMix).
|
|
|
|
Validates that the stock compiled extension runs rectangular kernels:
|
|
a (1, 9) / (9, 1) strip hits the existing K=9 template instantiation
|
|
(switch-K in dcnv4_im2col_cuda.cuh covers {9, 25, 49}; the kernel body is
|
|
generic over kernel_h/kernel_w at runtime).
|
|
|
|
Checks:
|
|
1. fwd correctness vs a pure-PyTorch bilinear reference for (1,9), (9,1)
|
|
and (3,3) sanity, fp32;
|
|
2. bwd correctness (grads wrt value and offset_mask) vs reference autograd;
|
|
3. padding channels of offset_mask (beyond G*K*3) are ignored;
|
|
4. fp16 fwd: finite, close to fp32 reference;
|
|
5. benchmark: strip (1,9) vs square 7 (K=49) vs square 3 (K=9) at the
|
|
SOFIA Stage-1 shape [8, 64*64, 192];
|
|
6. DCNv4Strip module smoke (zero-init => zero output; random init finite).
|
|
|
|
PASS gate: max|d| < 1e-4 fp32 (fwd and grads), fp16 finite, strip fwd+bwd
|
|
speedup >= 2.5x vs square 7.
|
|
|
|
Run on the training server (CUDA build of the extension required):
|
|
python scripts/test_dcnv4_strip.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
import torch
|
|
|
|
from DCNv4 import DCNv4Strip
|
|
from DCNv4.functions import DCNv4Function
|
|
|
|
ATOL_F32 = 1e-4
|
|
RTOL_F16 = 3e-2
|
|
|
|
|
|
def ref_dcnv4(
|
|
value: torch.Tensor, # [N, H, W, G*D]
|
|
offset_mask: torch.Tensor, # [N, H, W, P] (P >= G*K*3, tail = padding)
|
|
kh: int, kw: int,
|
|
pad_h: int, pad_w: int,
|
|
group: int,
|
|
offset_scale: float = 1.0,
|
|
) -> torch.Tensor:
|
|
"""Pure-PyTorch reference mirroring dcnv4_im2col_cuda.cuh semantics.
|
|
|
|
Point order: outer loop over i in [0, kw), inner over j in [0, kh)
|
|
(m = i*kh + j). Per group g the slab offset_mask[..., g*K*3:(g+1)*K*3]
|
|
holds K interleaved (dx, dy) pairs first, then K mask values (no softmax).
|
|
Sampling location for output pixel (hi, wi):
|
|
x = wi - pad_w + (i + dx) * offset_scale_adj, analogous for y,
|
|
which for stride=1, dilation=1 reduces to x = wi - pad_w + i + dx
|
|
when offset_scale == 1 (matches p0_w_/p0_h_ algebra in the kernel).
|
|
Bilinear sampling, zero outside the image. Differentiable.
|
|
"""
|
|
n, h, w, c = value.shape
|
|
d = c // group
|
|
k = kh * kw
|
|
dev = value.dtype
|
|
|
|
# Point grid in CUDA order: m = i*kh + j.
|
|
ii = torch.arange(kw, device=value.device).repeat_interleave(kh) # [K]
|
|
jj = torch.arange(kh, device=value.device).repeat(kw) # [K]
|
|
|
|
base_y = torch.arange(h, device=value.device).view(1, h, 1, 1)
|
|
base_x = torch.arange(w, device=value.device).view(1, 1, w, 1)
|
|
|
|
out = value.new_zeros(n, h, w, group, d)
|
|
half_w = (kw - 1) // 2
|
|
half_h = (kh - 1) // 2
|
|
for g in range(group):
|
|
slab = offset_mask[..., g * k * 3: (g + 1) * k * 3]
|
|
offs = slab[..., : k * 2].reshape(n, h, w, k, 2)
|
|
mask = slab[..., k * 2: k * 3] # [N,H,W,K]
|
|
dx, dy = offs[..., 0], offs[..., 1] # [N,H,W,K]
|
|
|
|
# p0 - centre*scale + (i*dil + dx)*scale (stride=1, dil=1)
|
|
x = (base_x - pad_w + half_w).to(dev) \
|
|
+ ((ii.view(1, 1, 1, k) - half_w).to(dev) + dx) * offset_scale
|
|
y = (base_y - pad_h + half_h).to(dev) \
|
|
+ ((jj.view(1, 1, 1, k) - half_h).to(dev) + dy) * offset_scale
|
|
|
|
x0 = torch.floor(x); y0 = torch.floor(y)
|
|
wx1 = x - x0; wy1 = y - y0
|
|
wx0 = 1.0 - wx1; wy0 = 1.0 - wy1
|
|
|
|
vg = value[..., g * d: (g + 1) * d] # [N,H,W,D]
|
|
acc = value.new_zeros(n, h, w, k, d)
|
|
for oy, wy_ in ((y0, wy0), (y0 + 1, wy1)):
|
|
for ox, wx_ in ((x0, wx0), (x0 + 1, wx1)):
|
|
inside = (oy >= 0) & (oy <= h - 1) & (ox >= 0) & (ox <= w - 1)
|
|
oy_c = oy.clamp(0, h - 1).long()
|
|
ox_c = ox.clamp(0, w - 1).long()
|
|
# Gather vg at [n, oy_c, ox_c] for every (h, w, k).
|
|
idx = (oy_c * w + ox_c).reshape(n, -1) # [N, H*W*K]
|
|
flat = vg.reshape(n, h * w, d)
|
|
samp = torch.gather(
|
|
flat, 1, idx.unsqueeze(-1).expand(-1, -1, d),
|
|
).reshape(n, h, w, k, d)
|
|
wgt = (wy_ * wx_ * inside.to(dev)).unsqueeze(-1)
|
|
acc = acc + samp * wgt
|
|
out[..., g, :] = (acc * mask.unsqueeze(-1)).sum(dim=3)
|
|
|
|
return out.reshape(n, h * w, c)
|
|
|
|
|
|
def run_op(value_seq, offset_mask, kh, kw, pad_h, pad_w, group, gc):
|
|
return DCNv4Function.apply(
|
|
value_seq.view(value_seq.shape[0], int(value_seq.shape[1] ** 0.5),
|
|
int(value_seq.shape[1] ** 0.5), -1),
|
|
offset_mask,
|
|
kh, kw, 1, 1, pad_h, pad_w, 1, 1,
|
|
group, gc, 1.0, 256, 0,
|
|
)
|
|
|
|
|
|
def check_case(kh: int, kw: int, n=2, h=16, w=16, c=64, group=4) -> None:
|
|
torch.manual_seed(42)
|
|
gc = c // group
|
|
k = kh * kw
|
|
p = ((group * k * 3 + 7) // 8) * 8
|
|
pad_h, pad_w = (kh - 1) // 2, (kw - 1) // 2
|
|
|
|
value = torch.randn(n, h, w, c, device="cuda", dtype=torch.float32)
|
|
om = torch.zeros(n, h, w, p, device="cuda", dtype=torch.float32)
|
|
om[..., : group * k * 3] = torch.randn(n, h, w, group * k * 3, device="cuda") * 0.7
|
|
|
|
v1 = value.clone().requires_grad_(True)
|
|
o1 = om.clone().requires_grad_(True)
|
|
v2 = value.clone().requires_grad_(True)
|
|
o2 = om.clone().requires_grad_(True)
|
|
|
|
out_op = DCNv4Function.apply(
|
|
v1, o1, kh, kw, 1, 1, pad_h, pad_w, 1, 1, group, gc, 1.0, 256, 0,
|
|
).reshape(n, h * w, c)
|
|
out_ref = ref_dcnv4(v2, o2, kh, kw, pad_h, pad_w, group)
|
|
|
|
d_fwd = (out_op - out_ref).abs().max().item()
|
|
assert d_fwd < ATOL_F32, f"({kh}x{kw}) fwd diverges: {d_fwd:.2e}"
|
|
|
|
g = torch.randn_like(out_op)
|
|
out_op.backward(g)
|
|
out_ref.backward(g)
|
|
d_gv = (v1.grad - v2.grad).abs().max().item()
|
|
d_go = (o1.grad[..., : group * k * 3] - o2.grad[..., : group * k * 3]).abs().max().item()
|
|
assert d_gv < ATOL_F32, f"({kh}x{kw}) grad_value diverges: {d_gv:.2e}"
|
|
assert d_go < ATOL_F32, f"({kh}x{kw}) grad_offset diverges: {d_go:.2e}"
|
|
|
|
# Padding channels must be ignored: poison them, output must not change.
|
|
om_poison = om.clone()
|
|
om_poison[..., group * k * 3:] = 1e6
|
|
out_poison = DCNv4Function.apply(
|
|
value, om_poison, kh, kw, 1, 1, pad_h, pad_w, 1, 1, group, gc, 1.0, 256, 0,
|
|
).reshape(n, h * w, c)
|
|
d_pad = (out_poison - out_op.detach()).abs().max().item()
|
|
assert d_pad == 0.0, f"({kh}x{kw}) padding channels are NOT ignored: {d_pad:.2e}"
|
|
|
|
# fp16 forward: finite + close to fp32 reference.
|
|
out_h = DCNv4Function.apply(
|
|
value.half(), om.half(), kh, kw, 1, 1, pad_h, pad_w, 1, 1,
|
|
group, gc, 1.0, 256, 0,
|
|
).reshape(n, h * w, c)
|
|
assert torch.isfinite(out_h).all(), f"({kh}x{kw}) fp16 produced non-finite"
|
|
rel = ((out_h.float() - out_ref).abs() / (out_ref.abs() + 1.0)).max().item()
|
|
assert rel < RTOL_F16, f"({kh}x{kw}) fp16 rel err {rel:.2e}"
|
|
|
|
print(f" OK ({kh}x{kw}): fwd {d_fwd:.1e}, dval {d_gv:.1e}, "
|
|
f"doff {d_go:.1e}, pad ignored, fp16 rel {rel:.1e}")
|
|
|
|
|
|
def bench(kh: int, kw: int, label: str, n=8, hw=64, c=192, group=12,
|
|
iters=50) -> float:
|
|
gc = c // group
|
|
k = kh * kw
|
|
p = ((group * k * 3 + 7) // 8) * 8
|
|
value = torch.randn(n, hw, hw, c, device="cuda", requires_grad=True)
|
|
om = torch.randn(n, hw, hw, p, device="cuda", requires_grad=True)
|
|
args = (kh, kw, 1, 1, (kh - 1) // 2, (kw - 1) // 2, 1, 1, group, gc, 1.0, 256, 0)
|
|
|
|
for _ in range(10): # warmup
|
|
out = DCNv4Function.apply(value, om, *args)
|
|
out.sum().backward()
|
|
torch.cuda.synchronize()
|
|
t0 = time.perf_counter()
|
|
for _ in range(iters):
|
|
out = DCNv4Function.apply(value, om, *args)
|
|
out.sum().backward()
|
|
torch.cuda.synchronize()
|
|
ms = (time.perf_counter() - t0) / iters * 1e3
|
|
print(f" {label:14s} K={k:2d}: {ms:7.3f} ms/iter (fwd+bwd, "
|
|
f"[{n},{hw}x{hw},{c}], G={group})")
|
|
return ms
|
|
|
|
|
|
def main() -> None:
|
|
assert torch.cuda.is_available(), "CUDA required"
|
|
print("== correctness ==")
|
|
check_case(3, 3) # sanity: known-good square path (K=9)
|
|
check_case(1, 9) # strip-H (K=9, reuses square-3 instantiation)
|
|
check_case(9, 1) # strip-V
|
|
check_case(1, 9, c=192, group=12) # SOFIA Stage-1 width
|
|
|
|
print("== DCNv4Strip module ==")
|
|
m = DCNv4Strip(channels=192, k=9, orientation="h", group=12).cuda()
|
|
assert (m.kernel_h, m.kernel_w) == (1, 9), "orientation='h' must give (1,9)"
|
|
mv = DCNv4Strip(channels=192, k=9, orientation="v", group=12).cuda()
|
|
assert (mv.kernel_h, mv.kernel_w) == (9, 1), "orientation='v' must give (9,1)"
|
|
x = torch.randn(2, 64 * 64, 192, device="cuda")
|
|
y = m(x, shape=(64, 64))
|
|
assert y.shape == x.shape
|
|
assert y.abs().max().item() < 1e-6, "zero-init must give ~zero output"
|
|
torch.nn.init.normal_(m.offset_mask.weight, std=0.02)
|
|
torch.nn.init.normal_(m.offset_mask.bias, std=0.02)
|
|
x = x.requires_grad_(True)
|
|
y = m(x, shape=(64, 64))
|
|
assert torch.isfinite(y).all() and y.abs().max().item() > 0
|
|
y.sum().backward()
|
|
assert x.grad is not None and torch.isfinite(x.grad).all()
|
|
print(" OK: module smoke (orientations, zero-init ~0, random init finite, bwd)")
|
|
|
|
print("== benchmark (SOFIA Stage-1 shape) ==")
|
|
t_strip = bench(1, 9, "strip (1,9)")
|
|
t_sq3 = bench(3, 3, "square 3")
|
|
t_sq7 = bench(7, 7, "square 7")
|
|
speedup = t_sq7 / t_strip
|
|
print(f" strip vs square7 fwd+bwd speedup: {speedup:.2f}x "
|
|
f"(target >= 2.5x, hard gate >= 1.8x), vs square3: {t_sq3 / t_strip:.2f}x")
|
|
if speedup < 2.5:
|
|
print(f" WARNING: speedup {speedup:.2f}x below 2.5x target "
|
|
f"(noisy benchmark? rerun on idle GPU)")
|
|
assert speedup >= 1.8, "GATE FAILED: strip speedup below hard gate 1.8x"
|
|
|
|
print("\nALL GATES PASSED — strip-DCN is viable on the stock binary.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|