File size: 4,952 Bytes
ed66dc1 | 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 | """torchvision.ops.deform_conv2d (modulated, stride 1, dilation 1, one offset group) in plain tensor ops that the LiteRT
converter lowers (GATHER on the token axis, elementwise bilinear weights, one matmul), for a CPU graph of BiRefNet's
decoder: the GPU has no gather, so the dense deformable convolution is a GPU wall (memory project_birefnet_park).
Same arithmetic as torchvision: for output pixel (h, w) and kernel tap k = (i, j), sample the input at
(h - pad + i + dy_k, w - pad + j + dx_k) bilinearly, corners outside the image count as 0, multiply by mask_k, then
contract over (tap, channel) with the weight."""
import torch
import torch.nn.functional as F
def deform_conv2d_gather(x, offset, weight, bias, padding, mask, max_bytes=None, lean=False, row0=0):
"""max_bytes: run the output rows in tiles so that no gathered [taps, rows, width, channels] float32 tensor passes
it (the CPU's peak memory; each output element keeps the same arithmetic). None: one tile.
lean: the corner-validity mask in min / max arithmetic (exact: the corner coordinates are integers) and the row
lookup as an embedding lookup (EMBEDDING_LOOKUP, no bounds-check ops: the indices are clamped), so the CPU runs
two ops per corner outside XNNPACK (the cast to int and the lookup) instead of about ten.
row0: offset and mask hold the output rows [row0, row0 + their height) only (a row tile of the whole output; the
input x is the whole map)."""
b, c, h, w = x.shape
o, _, kh, kw = weight.shape
k = kh * kw
ph, pw = (padding, padding) if isinstance(padding, int) else padding
ho, wo = offset.shape[2], w + 2 * pw - kw + 1
dev = x.device
# base sampling grid per tap: [k, ho, wo]
gy = torch.arange(ho, device=dev, dtype=x.dtype).view(1, ho, 1) + (row0 - ph)
gx = torch.arange(wo, device=dev, dtype=x.dtype).view(1, 1, wo) - pw
ki = torch.arange(kh, device=dev, dtype=x.dtype).repeat_interleave(kw).view(k, 1, 1)
kj = torch.arange(kw, device=dev, dtype=x.dtype).repeat(kh).view(k, 1, 1)
off = offset.view(b, k, 2, ho, wo)
msk = mask.view(b, k, ho, wo)
tokens = x.permute(0, 2, 3, 1).reshape(b * h * w, c) # token-major input: row = (batch, y, x)
base = (torch.arange(b, device=dev) * (h * w)).view(b, 1, 1, 1)
if lean:
base = base.to(x.dtype)
wmat = weight.permute(2, 3, 1, 0).reshape(k * c, o) # weight[o, c, i, j] -> [(i, j, c), o]
rows = ho if max_bytes is None else max(1, min(ho, int(max_bytes) // (b * k * wo * c * 4)))
outs = []
for r0 in range(0, ho, rows):
r1 = min(ho, r0 + rows)
n = r1 - r0
py = gy[:, r0:r1] + ki + off[:, :, 0, r0:r1] # [b, k, n, wo]
px = gx + kj + off[:, :, 1, r0:r1]
y0 = torch.floor(py)
x0 = torch.floor(px)
ly, lx = py - y0, px - x0
acc = 0
for dy, dx, wgt in ((0, 0, (1 - ly) * (1 - lx)), (0, 1, (1 - ly) * lx), (1, 0, ly * (1 - lx)), (1, 1, ly * lx)):
yy, xx = y0 + dy, x0 + dx
if lean:
# 1 inside, 0 outside; relu(t) - relu(t - 1) = clamp(t, 0, 1) without RELU_0_TO_1 (no XNNPACK kernel)
u = lambda t: F.relu(t) - F.relu(t - 1) # noqa: E731
valid = u(yy + 1) * u(h - yy) * u(xx + 1) * u(w - xx)
idx = (yy.clamp(0, h - 1) * w + xx.clamp(0, w - 1) + base).to(torch.int64) # [b, k, n, wo]
v = F.embedding(idx.reshape(-1), tokens).view(b, k, n, wo, c)
else:
valid = ((yy >= 0) & (yy <= h - 1) & (xx >= 0) & (xx <= w - 1)).to(x.dtype)
idx = (yy.clamp(0, h - 1) * w + xx.clamp(0, w - 1)).to(torch.int64) + base # [b, k, n, wo]
v = torch.index_select(tokens, 0, idx.reshape(-1)).view(b, k, n, wo, c)
acc = acc + v * (wgt * valid * msk[:, :, r0:r1]).unsqueeze(-1)
cols = acc.permute(0, 2, 3, 1, 4).reshape(b * n * wo, k * c) # (tap, channel) per output pixel
outs.append((cols @ wmat).view(b, n, wo, o))
out = (outs[0] if len(outs) == 1 else torch.cat(outs, dim=1)).permute(0, 3, 1, 2)
if bias is not None:
out = out + bias.view(1, o, 1, 1)
return out
if __name__ == '__main__':
from torchvision.ops import deform_conv2d
torch.manual_seed(0)
for ks, pad in ((1, 0), (3, 1), (7, 3)):
x = torch.randn(1, 64, 33, 40)
off = torch.randn(1, 2 * ks * ks, 33, 40) * 4
m = torch.rand(1, ks * ks, 33, 40) * 2
wt = torch.randn(256, 64, ks, ks) * 0.05
ref = deform_conv2d(x, off, wt, None, padding=pad, mask=m)
got = deform_conv2d_gather(x, off, wt, None, pad, m)
tiled = deform_conv2d_gather(x, off, wt, None, pad, m, max_bytes=ks * ks * 40 * 64 * 4 * 5, lean=True)
print(f'k={ks}: max |diff| {(ref - got).abs().max():.3e} tiled lean vs one tile {(tiled - got).abs().max():.3e} '
f'max |ref| {ref.abs().max():.3e}')
|