onnxruntime cpuid_info warning: Unknown CPU vendor. cpuinfo_vendor value: 15
/root/miniforge3/envs/mmcv_Birefnet_dcnv2/lib/python3.11/site-packages/mmcv/__init__.py:20: UserWarning: On January 1, 2023, MMCV will release v2.0.0, in which it will remove components related to the training process and add a data transformation module. In addition, it will rename the package names mmcv to mmcv-lite and mmcv-full to mmcv. See https://github.com/open-mmlab/mmcv/blob/master/docs/en/compatibility.md for more details.
  warnings.warn(
✅ Successfully registered mmcv::MMCVModulatedDeformConv2d symbolic function
✅ Model weights loaded successfully
✅ Model set to eval mode with im2col_step=1
🚀 Exporting ONNX model with fixed shape:
   - Input shape: (1, 3, 1024, 1024)
   - Output shape: [1, 1, 1024, 1024] (inferred)
/home/BiRefNet/models/backbones/swin_v1.py:441: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if W % self.patch_size[1] != 0:
/home/BiRefNet/models/backbones/swin_v1.py:443: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if H % self.patch_size[0] != 0:
/home/BiRefNet/models/backbones/swin_v1.py:382: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.
  Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size
/home/BiRefNet/models/backbones/swin_v1.py:382: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size
/home/BiRefNet/models/backbones/swin_v1.py:383: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.
  Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size
/home/BiRefNet/models/backbones/swin_v1.py:383: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size
/home/BiRefNet/models/backbones/swin_v1.py:213: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert L == H * W, "input feature has wrong size"
/home/BiRefNet/models/backbones/swin_v1.py:122: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert N == self.window_size[0] * self.window_size[1], "N must equal Wh*Ww for Swin window attention"
/home/BiRefNet/models/backbones/swin_v1.py:59: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  C = int(windows.shape[-1])
/home/BiRefNet/models/backbones/swin_v1.py:256: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if pad_r > 0 or pad_b > 0:
/home/BiRefNet/models/backbones/swin_v1.py:289: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert L == H * W, "input feature has wrong size"
/home/BiRefNet/models/backbones/swin_v1.py:294: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  pad_input = (H % 2 == 1) or (W % 2 == 1)
/home/BiRefNet/models/backbones/swin_v1.py:295: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if pad_input:
/root/miniforge3/envs/mmcv_Birefnet_dcnv2/lib/python3.11/site-packages/mmcv/ops/modulated_deform_conv.py:220: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if not all(map(lambda s: s > 0, output_size)):
❌ ONNX export failed: Argument passed to at() was not in the map.
Traceback (most recent call last):
  File "/home/BiRefNet/./export_onnx.py", line 130, in export_onnx
    torch.onnx.export(
  File "/root/miniforge3/envs/mmcv_Birefnet_dcnv2/lib/python3.11/site-packages/torch/onnx/utils.py", line 516, in export
    _export(
  File "/root/miniforge3/envs/mmcv_Birefnet_dcnv2/lib/python3.11/site-packages/torch/onnx/utils.py", line 1596, in _export
    graph, params_dict, torch_out = _model_to_graph(
                                    ^^^^^^^^^^^^^^^^
  File "/root/miniforge3/envs/mmcv_Birefnet_dcnv2/lib/python3.11/site-packages/torch/onnx/utils.py", line 1139, in _model_to_graph
    graph = _optimize_graph(
            ^^^^^^^^^^^^^^^^
  File "/root/miniforge3/envs/mmcv_Birefnet_dcnv2/lib/python3.11/site-packages/torch/onnx/utils.py", line 619, in _optimize_graph
    _C._jit_pass_peephole(graph, True)
IndexError: Argument passed to at() was not in the map.
❌ ONNX export failed. Please check the error messages above.

import torch
import torch.onnx
import numpy as np
from torch.onnx import symbolic_helper as sym_help
from utils import check_state_dict
from models.birefnet import BiRefNet
import warnings

# 忽略非关键警告
warnings.filterwarnings('ignore', category=UserWarning, module='torch.onnx')

# ==================== symbolic函数 ====================
def symbolic_MMCVModulatedDeformConv2d(g, input, offset, mask, weight, bias, stride, padding, dilation, groups, deform_groups):
    """
    正确的Symbolic function for mmcv::MMCVModulatedDeformConv2d
    
    参数必须与MMCV源码中的定义完全一致
    """
    # 处理bias为None的情况
    if bias is None or sym_help._is_none(bias):
        # 创建零偏置张量
        out_channels = weight.type().sizes()[0]
        bias = g.op("Constant", value_t=torch.zeros(out_channels, dtype=torch.float32))
    
    # 确保stride, padding, dilation是整数列表
    stride = sym_help._maybe_get_const(stride, 'ints')
    padding = sym_help._maybe_get_const(padding, 'ints')
    dilation = sym_help._maybe_get_const(dilation, 'ints')
    
    if not isinstance(stride, (list, tuple)):
        stride = [stride, stride]
    if not isinstance(padding, (list, tuple)):
        padding = [padding, padding]
    if not isinstance(dilation, (list, tuple)):
        dilation = [dilation, dilation]
    
    # 创建正确的ONNX节点
    return g.op(
        "mmcv::MMCVModulatedDeformConv2d",
        input,
        offset,
        mask,
        weight,
        bias,
        stride_i=stride,
        padding_i=padding,
        dilation_i=dilation,
        groups_i=groups,
        deform_groups_i=deform_groups
    )

# ==================== 注册自定义算子 ====================
# 关键修复：取消注释手动注册，因为register_extra_symbolics不包含我们需要的算子
torch.onnx.register_custom_op_symbolic(
    "mmcv::MMCVModulatedDeformConv2d", 
    symbolic_MMCVModulatedDeformConv2d, 
    opset_version=16
)

print("✅ Successfully registered mmcv::MMCVModulatedDeformConv2d symbolic function")

# ==================== 加载模型和权重 ====================
weights_file = "/home/BiRefNet/weights/BiRefNet-general-bb_swin_v1_tiny-epoch_232.pth"
device = "cpu"  # ONNX导出使用CPU

# 创建模型
birefnet = BiRefNet(bb_pretrained=False)

# 加载权重
state_dict = torch.load(weights_file, map_location="cpu")
state_dict = check_state_dict(state_dict)

# 权重映射 (DCNv2 适配)
new_state_dict = {}
for old_key, value in state_dict.items():
    # 映射 deform_conv.weight -> deform_conv.weight (直接使用)
    if 'deform_conv.weight' in old_key:
        new_state_dict[old_key] = value
    # 映射 bias 到 conv_bias
    elif 'deform_conv.bias' in old_key:
        new_key = old_key.replace('deform_conv.bias', 'conv_bias')
        new_state_dict[new_key] = value
    # 其他权重直接复制
    else:
        new_state_dict[old_key] = value

# 加载权重
birefnet.load_state_dict(new_state_dict, strict=False)
print("✅ Model weights loaded successfully")

# ==================== 设置 im2col_step=1 (batch_size=1) ====================
def set_im2col_step_for_batch1(model):
    """为 batch_size=1 设置所有 DCN 层的 im2col_step=1"""
    for name, module in model.named_modules():
        if hasattr(module, 'im2col_step'):
            module.im2col_step = 1
            print(f"🔧 Setting im2col_step=1 for layer: {name}")
        # 特别处理 BiRefNet 中的 DCN 层
        if hasattr(module, 'deform_conv') and hasattr(module.deform_conv, 'im2col_step'):
            module.deform_conv.im2col_step = 1
            print(f"🔧 Setting im2col_step=1 for {name}.deform_conv")
    
    return model

birefnet = set_im2col_step_for_batch1(birefnet)

# 设置为 eval 模式
birefnet.eval()
print("✅ Model set to eval mode with im2col_step=1")

# ==================== 固定形状的ONNX导出函数 ====================
def export_onnx(net, output_path, input_shape=(1, 3, 1024, 1024)):
    """
    固定形状的ONNX导出函数
    
    参数:
        net: 模型
        output_path: 输出路径
        input_shape: 固定的输入形状
    """
    # 创建固定形状的示例输入
    input_tensor = torch.randn(*input_shape)
    
    print(f"🚀 Exporting ONNX model with fixed shape:")
    print(f"   - Input shape: {input_shape}")
    print(f"   - Output shape: [1, 1, 1024, 1024] (inferred)")
    
    try:
        # 导出ONNX模型 - 固定形状
        torch.onnx.export(
            net,
            input_tensor,
            output_path,
            opset_version=16,
            input_names=['input_image'],
            output_names=['output_image'],
            export_params=True,
            do_constant_folding=True,
            verbose=False,
            training=torch.onnx.TrainingMode.EVAL,
        )
        
        print(f"✅ ONNX model exported successfully to {output_path}")
        return True
        
    except Exception as e:
        print(f"❌ ONNX export failed: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

# ==================== 执行导出 ====================
output_onnx_file = weights_file.replace('.pth', '_dcnv2.onnx')

# 执行导出
success = export_onnx(birefnet, output_onnx_file)

if success:
    print("🎉 ONNX export completed successfully!")
    # 验证导出的ONNX模型
    try:
        import onnx
        onnx_model = onnx.load(output_onnx_file)
        onnx.checker.check_model(onnx_model)
        print("✅ ONNX model validation passed!")
        
        # 打印模型输入输出信息
        print(f"Model input shape: {[dim.dim_value for dim in onnx_model.graph.input[0].type.tensor_type.shape.dim]}")
        print(f"Model output shape: {[dim.dim_value for dim in onnx_model.graph.output[0].type.tensor_type.shape.dim]}")
        
    except Exception as e:
        print(f"⚠️ ONNX model validation warning: {str(e)}")
else:
    print("❌ ONNX export failed. Please check the error messages above.")

https://github.com/ZhengPeng7/BiRefNet