| |
| |
| |
| 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 |
| 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) |
|
|
