如何在Node.js中实现Restormer图像复原Transformer模型？

Restormer是CVPR 2022上提出的高分辨率图像复原Transformer模型，在去噪、去雨、去模糊、低光照增强等多个任务上都拿到了领先成绩。它的原始实现基于PyTorch，很多前端和全栈开发者的技术栈里并没有Python环境，那么能不能直接在Node.js里跑起来呢？答案是肯定的。通过把PyTorch训练好的权重导出为ONNX格式，再用onnxruntime-node加载推理，我们可以在纯JavaScript技术栈中完成完整的复原流水线。这篇文章就从原理到落地，把整个实现过程讲清楚。

Restormer的核心原理：为什么它能又快又准

传统的Vision Transformer在处理图像时，自注意力机制需要在空间维度上计算相似度矩阵。假设特征图尺寸为C乘以H乘以W，标准注意力的时间复杂度是关于HW的平方级。对于图像复原这种必须处理高分辨率输入的任务来说，这个开销几乎是不可接受的，一张1024乘以1024的图像直接算下来显存就爆了。

Restormer的关键创新是多头转置注意力（Multi-Dconv Head Transposed Cross-Attention，MDTA）。它把注意力的计算维度从空间转到了通道：先在通道维度上计算全局相关性，得到的是C乘以C的相似度矩阵，而C通常是远小于HW的。这样一来复杂度变成了关于通道数的平方、关于像素数的线性，高分辨率图像也能吃得下。同时在注意力之前加入了深度可分离卷积，用来注入局部上下文信息，弥补通道注意力对空间结构感知不足的问题。

另一个核心模块是门控Dconv前馈网络（Gated-Dconv Feed-Forward Network，GDFN）。普通Transformer的前馈网络就是两个线性层加激活，GDFN则把特征分成两路，一路过门控控制信息流，另一路保留内容信息，再用逐点卷积融合。这种双分支设计让网络在去噪时能够选择性地保留干净的结构、抑制噪声相关的特征，实验证明对复原质量的提升非常明显。

整体架构上，Restormer采用多尺度编解码结构，类似U-Net：从输入层开始逐级下采样提取语义，再逐级上采样恢复分辨率，最后用一个精细化阶段做细节重建。训练时采用渐进式策略，先在小块上训练，再逐步增大图像尺寸，这让它对高分辨率推理的泛化性特别好。

模型导出：从PyTorch权重到ONNX文件

要在Node.js里使用Restormer，第一步是拿到ONNX模型。官方仓库提供了各任务的预训练权重，我们可以用PyTorch加载后导出。导出时要注意动态维度的问题：如果希望模型支持任意尺寸输入，需要把动态轴设为dynamic_axes；如果只处理固定尺寸，导出静态形状的模型推理速度会更快，onnxruntime也更容易做图优化。

import torch
from basicsr.models.archs.restormer_arch import Restormer

# 以去雨任务为例加载预训练权重
model = Restormer(
    inp_channels=3, out_channels=3, dim=48,
    n_blocks=4, n_refinement_blocks=4,
    heads=[1, 2, 4, 8], ffn_expansion_factor=2.66,
    bias=False, LayerNorm_type='BiasFree'
)
ckpt = torch.load('Deraining.pth', map_location='cpu')
model.load_state_dict(ckpt['params'])
model.eval()

dummy = torch.randn(1, 3, 256, 256)
torch.onnx.export(
    model, dummy, 'restormer_derain.onnx',
    input_names=['input'], output_names=['output'],
    opset_version=17,
    dynamic_axes={
        'input': {0: 'batch', 2: 'height', 3: 'width'},
        'output': {0: 'batch', 2: 'height', 3: 'width'}
    }
)

导出完成后建议用onnx.checker校验一下模型文件，再用onnxruntime的Python包跑一次对比输出，确保导出过程中的数值误差在可接受范围内。LayerNorm这类算子在低版本opset下容易出兼容问题，建议opset版本设为17以上，可以省去很多麻烦。

还有一个实用技巧：如果部署机器的CPU支持AVX2或AVX512，可以在导出后做一次FP16量化测试，很多情况下精度损失几乎为零，但推理耗时能下降百分之三十左右。不过FP16主要收益在GPU上，纯CPU场景更推荐保持FP32或者做INT8量化。

Node.js推理实现：完整代码落地

环境准备很简单，安装onnxruntime-node和一个图像处理库。图像解码推荐用sharp，性能比纯JS方案好得多，而且能直接输出模型需要的RGB浮点缓冲区。

npm init -y
npm install onnxruntime-node sharp

下面是完整的推理脚本。流程分为四步：读取图片并resize到模型输入尺寸、归一化到0到1区间并调整布局为NCHW、执行session推理、把输出还原成图像写回磁盘。Restormer的输入是0到1的RGB值，不需要做均值方差标准化，这一点和分类模型不同，写错了会直接导致输出发灰或者出现色块。

const ort = require('onnxruntime-node');
const sharp = require('sharp');

async function main() {
  const session = await ort.InferenceSession.create('restormer_derain.onnx', {
    executionProviders: ['cpu'],
    graphOptimizationLevel: 'all'
  });

  // 读取图像并统一到256x256，与导出时的dummy输入一致
  const img = sharp('rainy_input.jpg');
  const { data, info } = await img
    .resize(256, 256, { fit: 'fill' })
    .removeAlpha()
    .raw()
    .toBuffer({ resolveWithObject: true });

  // HWC转NCHW，像素值归一化到0-1
  const [h, w] = [info.height, info.width];
  const chw = new Float32Array(3 * h * w);
  for (let i = 0; i < h * w; i++) {
    for (let c = 0; c < 3; c++) {
      chw[c * h * w + i] = data[i * info.channels + c] / 255.0;
    }
  }

  const input = new ort.Tensor('float32', chw, [1, 3, h, w]);
  const results = await session.run({ input });
  const output = results.output.data;

  // NCHW转回HWC并写图
  const outBuf = Buffer.alloc(h * w * 3);
  for (let i = 0; i < h * w; i++) {
    for (let c = 0; c < 3; c++) {
      outBuf[i * 3 + c] = Math.max(0, Math.min(255, output[c * h * w + i] * 255));
    }
  }
  await sharp(outBuf, { raw: { width: w, height: h, channels: 3 } })
    .jpeg({ quality: 95 })
    .toFile('restored_output.jpg');
  console.log('图像复原完成');
}

main().catch(console.error);

代码中有两个容易踩坑的地方值得展开说。第一是通道布局的转换，ONNX模型默认NCHW格式，而sharp输出的raw buffer是HWC排列，两层循环转换时索引千万别写反，可以用小尺寸图像先验证一遍。第二是动态尺寸问题，如果导出时启用了dynamic_axes，推理时理论上可以传任意尺寸，但Restormer内部有多层下采样，输入宽高最好是2的幂次或者至少能被16整除，否则可能出现尺寸对不上的报错。

对于超过显存或内存的大图，推荐做分块推理：把图像切成若干个重叠的小块分别复原，再在重叠区域做线性融合。块间重叠取32到64像素通常就能避免明显的拼接痕迹。同时建议把推理部分封装成异步任务队列，Node.js的事件循环天然适合这种IO密集加计算密集混合的场景，配合worker_threads可以把模型推理放到独立线程，避免阻塞HTTP服务的响应。

性能分析与部署建议

在一台普通四核CPU机器上实测，256乘256尺寸的单张推理大约在几百毫秒到一秒出头，尺寸翻倍耗时大约涨到三到四倍，基本符合像素数线性增长的预期。如果需要服务化部署，有几个优化方向：启用onnxruntime的intraOpNumThreads并把线程数设为物理核数；对CPU部署开启图优化级别的all选项；批量请求时用batch维度合并多张图一次推理，吞吐量提升明显。

内存方面要注意Tensor的生命周期。onnxruntime-node的Tensor对象在JS侧被垃圾回收前会一直持有底层缓冲区，长时间运行的服务如果频繁创建大Tensor，建议在推理完成后主动把引用置空，并观察进程RSS是否稳定。另外session实例应当全局复用，每次请求都create一个session会重复加载几十MB的模型文件，这是新手最常犯的错误之一。

如果对延迟要求更高，可以考虑切到GPU执行：onnxruntime-node支持CUDA execution provider，需要额外安装CUDA和cuDNN的动态库，安装目录建议通过环境变量指定。Node.js版本建议使用18以上的LTS，onnxruntime-node对旧版本Node的ABI兼容性并不完美，遇到加载失败的报错可以先检查版本匹配情况。

整体来说，Restormer配合ONNX Runtime是一套非常适合全栈JavaScript团队的图像复原方案：模型精度有保障，部署链路完全不需要引入Python运行时，从图片上传接口到复原输出可以在同一个服务里闭环完成。只要处理好输入预处理、分块策略和线程调度这几个细节，在生产环境中跑起来并不困难。

Node.jsRestormerTransformer图像复原修改时间：2026-09-09 19:32:53