U2Net模型.pth转.onnx后性能下降，背景分割效果模糊

我要给网站添加汽车轮毂图像的背景移除功能，选用rembg库实现，该库依赖U2Net模型。标准U2Net只能移除轮毂外部背景，辐条、孔洞等内部区域无法处理，因此我对U2Net做了针对性微调。

微调后的.pth模型生成的黑白分割Mask完全符合需求，但将模型转换为rembg所需的.onnx格式后，生成的Mask变得模糊、边缘虚化。即使转换未微调的标准U2Net模型，结果同样模糊，背景分割功能失效，说明问题出在模型转换环节。

效果示例 #

• 原始图像：原始图像
• 微调后.pth模型生成的Mask：微调后PTH模型生成的Mask
• 转换为.onnx后生成的Mask：转ONNX后生成的Mask

生成Mask的命令 #

在U2Net中执行： #

python3 u2net_test.py

在rembg中执行： #

rembg i -om -m u2net_custom -x '{"model_path": "~/.u2net/u2net_custom.onnx"}' 55.jpg 55.png

尝试的两种模型转换方式 #

1. 转换已训练完成的模型 #

import torch
import torch.onnx
from model.u2net import U2NET

def load_model(model_path, model_class):
    checkpoint = torch.load(model_path, map_location='cpu')
    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
        model = model_class()
        model.load_state_dict(checkpoint['state_dict'])
    else:
        model = model_class()
    model.eval()
    return model

def convert_to_onnx(model, output_path):
    dummy_input = torch.randn(1, 3, 320, 320)
    torch.onnx.export(model, dummy_input, output_path, opset_version=12,
                          dynamic_axes={'input': {0: 'batch_size', 2: 'height', 3: 'width'},
                                        'output': {0: 'batch_size', 2: 'height', 3: 'width'}})
    print(f"success {output_path}")

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="conversion PyTorch to ONNX")
    parser.add_argument('--model-path', type=str, required=True, help='path to .pth file')
    parser.add_argument('--output-path', type=str, required=True, help='save ONNX file')

    args = parser.parse_args()
    model = load_model(args.model_path, U2NET)
    convert_to_onnx(model, args.output_path)

2. 训练过程中同步保存ONNX模型 #

if ite_num % save_frq == 0:
    timestamp = int(time.time())
    filePath = model_dir + model_name+"_%d_%d." % (ite_num, timestamp)

    torch.save(net.state_dict(), filePath + 'pth')

    dummy_input = torch.randn(1, 3, 320, 320)
    net.eval()
    torch.onnx.export(net, dummy_input, filePath + 'onnx', opset_version=12)

    running_loss = 0.0
    running_tar_loss = 0.0
    net.train()  # 恢复训练
    ite_num4val = 0

我尝试过调整转换设置、更换库版本、修改opset_version等方法，结果均一致：转换后的模型完全失效。请问问题出在哪里？

内容的提问来源于stack exchange，提问作者Yuriy