3步搞定U-2-Net模型导出ONNX：从PyTorch到跨平台部署全指南

你是否在模型部署时遇到框架限制？U-2-Net作为优秀的显著对象检测模型，原生PyTorch格式难以直接用于生产环境。本文将通过3个核心步骤，教你如何将U-2-Net模型转换为ONNX（开放神经网络交换格式），实现跨框架部署。完成后，你将掌握模型加载、转换优化和验证的全流程，解决90%的部署兼容性问题。

准备工作：环境与模型文件

环境依赖安装

首先确保已安装必要依赖。项目依赖清单见requirements.txt，关键包包括torch、onnx和onnxruntime。使用以下命令安装：

pip install -r requirements.txt
pip install onnx onnxruntime

模型文件获取

U-2-Net提供两种模型：基础版和人像专用版。通过setup_model_weights.py脚本自动下载权重：

python setup_model_weights.py

脚本会将权重保存到：

• 基础模型：saved_models/u2net/u2net.pth
• 人像模型：saved_models/u2net_portrait/u2net_portrait.pth

模型网络结构定义在model/u2net.py，包含U2NET和U2NETP（轻量级）两个类，核心采用嵌套U型结构设计。

步骤1：加载预训练模型

创建转换脚本export_onnx.py，首先加载PyTorch模型。以下代码片段展示如何加载基础模型：

import torch
from model.u2net import U2NET

# 初始化模型
model = U2NET(3, 1)  # 3通道输入，1通道输出
device = torch.device('cpu')
model.to(device)

# 加载权重
model.load_state_dict(torch.load(
    'saved_models/u2net/u2net.pth', 
    map_location=device
))
model.eval()  # 设置为推理模式

关键注意事项：

• 使用map_location=device确保CPU加载，避免GPU依赖
• 必须调用model.eval()禁用dropout和批量归一化（Batch Normalization）的训练模式行为
• 人像模型加载只需将U2NET替换为U2NETP，权重路径改为对应人像模型路径

步骤2：导出ONNX格式

基本导出代码

添加以下代码到export_onnx.py完成转换：

# 创建输入张量（batch_size=1, 3通道, 320x320分辨率）
input_tensor = torch.randn(1, 3, 320, 320, device=device)

# 导出ONNX
torch.onnx.export(
    model,
    input_tensor,
    'u2net.onnx',  # 输出路径
    opset_version=11,  # 操作集版本
    do_constant_folding=True,  # 优化常量折叠
    input_names=['input'],  # 输入节点名称
    output_names=['output'],  # 输出节点名称
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size', 2: 'height', 3: 'width'}
    }  # 动态维度支持
)

参数优化说明

• opset_version=11：选择ONNX 11版本以兼容主流推理引擎
• dynamic_axes：设置动态批次和分辨率，支持不同输入尺寸
• 对于人像模型，输出文件可命名为u2net_portrait.onnx

导出轻量级模型

若需导出轻量级模型U2NETP，只需修改模型初始化：

from model.u2net import U2NETP
model = U2NETP(3, 1)  # 轻量级版本

步骤3：验证ONNX模型

推理结果对比

使用ONNX Runtime加载转换后的模型，与PyTorch输出对比：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
ort_session = ort.InferenceSession('u2net.onnx')
input_name = ort_session.get_inputs()[0].name

# PyTorch推理
with torch.no_grad():
    torch_output = model(input_tensor)

# ONNX推理
onnx_output = ort_session.run(
    None,
    {input_name: input_tensor.numpy()}
)

# 验证误差（应小于1e-5）
np.testing.assert_allclose(
    torch_output[0].numpy(), 
    onnx_output[0], 
    rtol=1e-5, 
    atol=1e-5
)

可视化检测效果

U-2-Net擅长人像分割和背景移除任务。以下是模型处理效果示例：

人像分割结果：

背景移除对比：

部署应用场景

转换后的ONNX模型可部署到多种平台：

• 移动端：通过ONNX Runtime Mobile部署到iOS/Android
• Web端：使用ONNX.js在浏览器中运行
• 服务端：结合FastAPI构建推理服务

项目提供的Gradio演示gradio/demo.py可直接使用ONNX模型加速推理，界面示例：

常见问题解决

动态尺寸不支持

若部署环境要求固定输入尺寸，修改导出代码移除dynamic_axes参数，指定固定分辨率：

input_tensor = torch.randn(1, 3, 512, 512)  # 固定512x512

模型体积过大

基础模型转换后约170MB，可通过以下方式优化：

1. 使用轻量级模型U2NETP（约40MB）
2. 启用ONNX优化工具：

python -m onnxruntime.tools.optimize_onnx_model u2net.onnx --output u2net_opt.onnx

推理结果差异

若PyTorch与ONNX结果差距较大，检查：

• 是否调用model.eval()
• 输入数据预处理是否一致
• ONNX导出时指定opset_version=11以上

总结与后续步骤

本文通过3个步骤完成U-2-Net到ONNX的转换：

1. 环境准备与模型下载
2. PyTorch模型加载与ONNX导出
3. 模型验证与优化

下一步建议：

• 尝试人像专用模型转换：saved_models/u2net_portrait/u2net_portrait.pth
• 探索量化技术进一步减小模型体积：onnxruntime.quantization
• 参考test_data/中的示例图片进行批量推理测试

完整转换脚本export_onnx.py已添加到项目根目录，包含基础版和人像版两种转换选项。通过ONNX格式，U-2-Net模型可无缝集成到TensorRT、OpenVINO等推理框架，实现跨平台高效部署。

点赞收藏本文，关注项目README.md获取更多部署教程，下期将讲解如何使用ONNX Runtime部署到Android应用。