File size: 3,276 Bytes
05b7816 | 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 | # BiRefNet_lite-matting for onnxruntime-web's WebGPU EP: each deformable conv becomes one GridSample + MatMul (float math
# only), and the sampling grid's base is made at run time from a CumSum over the offsets (a data tensor), so constant
# folding can't bake a full-size grid per layer into the file (that made the fp16 file 150-177 MB).
import sys, torch, torch.nn.functional as F, onnx, onnxslim
from transformers import AutoModelForImageSegmentation
from onnxruntime.transformers.float16 import convert_float_to_float16
def deform_conv2d_rt(input, offset, weight, bias=None, stride=(1, 1), padding=(0, 0), dilation=(1, 1), mask=None):
n, c, h, w = input.shape
o, _, kh, kw = weight.shape
sh, sw = (stride, stride) if isinstance(stride, int) else stride
ph, pw = (padding, padding) if isinstance(padding, int) else padding
dh, dw = (dilation, dilation) if isinstance(dilation, int) else dilation
ho, wo = offset.shape[2], offset.shape[3]
k = kh * kw
ones = offset[:, :1] * 0 + 1 # n, 1, ho, wo
oy = torch.cumsum(ones, 2) - 1
ox = torch.cumsum(ones, 3) - 1
ky, kx = torch.meshgrid(torch.arange(kh, dtype=input.dtype), torch.arange(kw, dtype=input.dtype), indexing='ij')
ky = (ky.reshape(1, k, 1, 1) * dh)
kx = (kx.reshape(1, k, 1, 1) * dw)
off = offset.reshape(n, k, 2, ho, wo)
y = oy * sh - ph + ky + off[:, :, 0]
x = ox * sw - pw + kx + off[:, :, 1]
grid = torch.stack([x * (2 / (w - 1)) - 1, y * (2 / (h - 1)) - 1], -1).reshape(n, k * ho, wo, 2)
s = F.grid_sample(input, grid, mode='bilinear', padding_mode='zeros', align_corners=True).reshape(n, c, k, ho, wo)
if mask is not None:
s = s * mask.reshape(n, 1, k, ho, wo)
out = torch.matmul(weight.reshape(o, c * k), s.reshape(n, c * k, ho * wo)).reshape(n, o, ho, wo)
if bias is not None:
out = out + bias.reshape(1, o, 1, 1)
return out
model = AutoModelForImageSegmentation.from_pretrained('models/birefnet', trust_remote_code=True).eval()
mod = sys.modules[type(model).__module__]
orig = mod.deform_conv2d
x = torch.randn(1, 8, 20, 24); off = torch.randn(1, 18, 20, 24) * 2; m = torch.rand(1, 9, 20, 24) * 2
wt = torch.randn(5, 8, 3, 3); b = torch.randn(5)
print('op diff', float((orig(x, off, wt, b, padding=(1, 1), mask=m) - deform_conv2d_rt(x, off, wt, b, padding=(1, 1), mask=m)).abs().max()))
ims = torch.randn(1, 3, 512, 512)
with torch.no_grad():
a = model(ims)[-1]
mod.deform_conv2d = deform_conv2d_rt
b2 = model(ims)[-1]
print('model diff', float((a - b2).abs().max()))
for size in [int(s) for s in sys.argv[1:]] or [1024]:
raw = f'models/birefnet-rt-{size}.onnx'
with torch.no_grad():
torch.onnx.export(model, torch.randn(1, 3, size, size), raw, opset_version=17,
input_names=['input_image'], output_names=['output_image'], dynamo=False)
slim = onnxslim.slim(onnx.load(raw))
m16 = convert_float_to_float16(slim, keep_io_types=True)
out = f'models/birefnet-lite-matting-{size}-rt-fp16.onnx'
onnx.save(m16, out)
ops = {}
for nd in m16.graph.node: ops[nd.op_type] = ops.get(nd.op_type, 0) + 1
print('saved', out, len(m16.graph.node), 'nodes', sorted(ops.items(), key=lambda t: -t[1])[:12], flush=True)
|