如何在浏览器（客户端）运行ONNX格式U2Net模型并解决推理效果差问题

你遇到的输出异常主要由以下几个可复现的错误导致，按优先级修复即可：

• 通道顺序不匹配：Canvas获取的图像像素是HWC格式（按像素点顺序存储每个点的R、G、B、A值），而U2Net的ONNX模型要求输入为NCHW格式（先存储整图所有R通道数值，再存G通道，最后存B通道）。你当前的代码直接按R0、G0、B0、R1、G1、B1...的顺序填充张量，对应到NCHW格式会导致通道完全错位，这是输出效果极差的核心原因。
• 精度丢失问题：预处理时调用toFixed(2)会将像素值强制截断为两位小数，额外丢失图像信息，直接做除法运算即可无需该操作。
• 输出缺少激活处理：U2Net原生输出的是logits值，需要经过Sigmoid激活才能得到0~1范围的掩码概率，如果你导出ONNX时没有把Sigmoid算子封装到模型内，直接将原始输出乘255会得到完全错误的像素值。
• 语法逻辑问题：创建推理会话时await和then混用的写法不规范，可能导致会话初始化异常。

核心代码修复示例 #

预处理部分（修正通道顺序） #

const height = 320, width = 320;
const inputBuffer = new Float32Array(3 * height * width);
const imgData = input_imageData.data;
// 按CHW顺序填充张量
for (let c = 0; c < 3; c++) {
  for (let h = 0; h < height; h++) {
    for (let w = 0; w < width; w++) {
      const imgPos = (h * width + w) * 4 + c;
      const tensorPos = c * height * width + h * width + w;
      const normVal = imgData[imgPos] / 255;
      // 对应通道标准化
      if (c === 0) inputBuffer[tensorPos] = (normVal - 0.485) / 0.229;
      else if (c === 1) inputBuffer[tensorPos] = (normVal - 0.456) / 0.224;
      else inputBuffer[tensorPos] = (normVal - 0.406) / 0.225;
    }
  }
}
const input = new ort.Tensor('float32', inputBuffer, [1, 3, height, width]);

后处理部分（添加Sigmoid激活） #

const predData = pred.data;
const myImageData = ctx.createImageData(width, height);
for (let i = 0; i < predData.length; i++) {
  // Sigmoid激活将logits映射到0~1区间
  const maskVal = 1 / (1 + Math.exp(-predData[i]));
  const imgPos = i * 4;
  myImageData.data[imgPos] = Math.round(maskVal * 255);
  myImageData.data[imgPos + 1] = Math.round(maskVal * 255);
  myImageData.data[imgPos + 2] = Math.round(maskVal * 255);
  myImageData.data[imgPos + 3] = 255;
}
ctx.putImageData(myImageData, 0, 0);

附加验证步骤 #

修复后如果仍有偏差，可以按以下顺序排查：

1. 拿同一张测试图在Python端用相同的预处理逻辑跑PyTorch原生模型，对比JS端预处理后的张量数值，确认输入完全一致。
2. 打开ONNX模型确认输入节点名称确实为input.1，输出形状为[1,1,320,320]，避免节点配置不匹配。
3. 确认ONNX导出时opset版本≥12，未做错误的量化、剪枝操作。

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