Implementation:PeterL1n BackgroundMattingV2 Torch onnx export

Knowledge Sources	BackgroundMattingV2 PyTorch ONNX Export
Domains	Model_Deployment, Interoperability
Last Updated	2026-02-09 00:00 GMT

Overview

Concrete wrapper for torch.onnx.export as used in BackgroundMattingV2's ONNX export script with compatibility options for patch operations.

Description

The ONNX export in export_onnx.py wraps torch.onnx.export with configuration specific to the matting model. Key features:

Dynamic axes for batch, height, and width dimensions on all inputs and outputs
Compatibility options for the refiner's patch crop method (roi_align for ONNX, instead of default unfold) and replace method (scatter_element for ONNX, instead of default scatter_nd)
Dummy inputs at 1080×1920 resolution with batch size 2 for tracing
Named inputs/outputs (src, bgr → pha, fgr, etc.)
Precision support (float32 or float16)

Usage

Use when deploying the matting model to non-PyTorch runtimes (ONNX Runtime, TensorRT, OpenVINO). Choose ONNX-compatible patch methods for the refiner.

Code Reference

Source Location

Repository: BackgroundMattingV2
File: export_onnx.py
Lines: 112-121

Signature

torch.onnx.export(
    model=model,                    # MattingBase or MattingRefine
    args=(src, bgr),               # Dummy inputs: Tensor[2,3,1080,1920]
    f=args.output,                 # Output .onnx file path
    verbose=args.onnx_verbose,     # Print export details (default True)
    opset_version=args.onnx_opset_version,  # ONNX opset (default 12)
    do_constant_folding=args.onnx_constant_folding,  # Fold constants (default True)
    input_names=['src', 'bgr'],
    output_names=['pha', 'fgr', 'pha_sm', 'fgr_sm', 'err_sm', 'ref_sm'],  # MattingRefine
    dynamic_axes={name: {0: 'batch', 2: 'height', 3: 'width'} for name in all_names}
)

Import

import torch
from model import MattingBase, MattingRefine

I/O Contract

Inputs

Name	Type	Required	Description
model	nn.Module	Yes	MattingBase or MattingRefine in eval mode
model-checkpoint	str	Yes	Path to trained checkpoint
onnx-opset-version	int	No	ONNX opset version (default 12)
precision	str	No	'float32' or 'float16' (default 'float32')
model-refine-patch-crop-method	str	No	'roi_align' for ONNX compat (default)
model-refine-patch-replace-method	str	No	'scatter_element' for ONNX compat (default)
output	str	Yes	Output .onnx file path

Outputs

Name	Type	Description
.onnx file	File	ONNX model with dynamic axes for batch/height/width

Usage Examples

Export MattingRefine to ONNX

import torch
from model import MattingRefine

model = MattingRefine(
    backbone='resnet50',
    backbone_scale=0.25,
    refine_mode='sampling',
    refine_sample_pixels=80000,
    refine_patch_crop_method='roi_align',      # ONNX compatible
    refine_patch_replace_method='scatter_element'  # ONNX compatible
)
model.load_state_dict(torch.load('checkpoint.pth', map_location='cpu'), strict=False)
model.eval().float()

# Dummy inputs for tracing
src = torch.randn(2, 3, 1080, 1920)
bgr = torch.randn(2, 3, 1080, 1920)

torch.onnx.export(
    model, (src, bgr), 'matting.onnx',
    opset_version=12,
    input_names=['src', 'bgr'],
    output_names=['pha', 'fgr', 'pha_sm', 'fgr_sm', 'err_sm', 'ref_sm'],
    dynamic_axes={name: {0: 'batch', 2: 'height', 3: 'width'}
                  for name in ['src', 'bgr', 'pha', 'fgr', 'pha_sm', 'fgr_sm', 'err_sm', 'ref_sm']}
)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment