Swin2SR是2022年提出的一个基于Swin Transformer V2骨干网络的超分辨率重建模型，在图像修复和细节恢复任务上表现相当出色。大多数教程都基于Python环境，但如果你现有的技术栈是Node.js，比如Electron桌面应用或者后端图片处理服务，引入Python会带来额外的部署负担。其实借助ONNX Runtime的Node.js绑定，完全可以在纯Node环境中跑通Swin2SR推理，本文就完整演示这一过程。
准备工作：获取并转换Swin2SR模型
第一步是把PyTorch格式的Swin2SR权重转换成ONNX格式，因为Node.js端没有成熟的PyTorch运行时，而ONNX Runtime对Node的支持非常完善。转换工作只需要做一次，可以在任何一台装了Python的机器上完成，转换完成后Node环境就不再依赖Python了。
从官方仓库下载预训练权重，推荐先用Swin2SR_ClassicalSR_X2_64这一款，输入尺寸为64x64的小块，放大倍数为2倍，适合作为入门验证。转换脚本的核心是设置好dummy输入然后调用torch.onnx.export，示例代码如下：
import torch
from models import Swin2SR
model = Swin2SR(upscale=2, img_size=64, window_size=8,
img_range=1., depths=[6]*6, embed_dim=180,
num_heads=[6]*6, mlp_ratio=2, upsampler='pixelshuffledirect')
model.load_state_dict(torch.load('Swin2SR_ClassicalSR_X2_64.pth')['params_ema'], map_location='cpu')
model.eval()
dummy = torch.randn(1, 3, 64, 64)
torch.onnx.export(model, dummy, 'swin2sr_x2.onnx',
input_names=['input'],
output_names=['output'],
opset_version=16,
dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'},
'output': {0: 'batch', 2: 'height', 3: 'width'}})
print('导出完成')
注意这里启用了动态维度，这样推理时不再局限于64x64的输入，可以直接传入整张图。但要注意Swin2SR对输入尺寸有对齐要求，window_size为8时，宽高最好能被8整除，否则需要在Node端做padding预处理，推理完再裁剪回来。
Node.js端环境搭建与推理封装
Node端需要安装两个核心依赖：onnxruntime-node负责模型推理，sharp负责图像解码、缩放和通道转换。前者自带各平台预编译二进制，Windows、Linux、macOS都能直接使用，不需要额外配置CUDA（CPU推理足够可用，GPU加速后面单独说）。
npm init -y npm install onnxruntime-node sharp
封装推理时要处理的关键点是数据格式。ONNX模型的输入是一个四维浮点张量，形状为NCHW，取值范围0到1。而sharp解码出来的是RGB排列的Uint8数据，所以需要手动做两件事：把HWC转成CHW，把整数除以255做归一化。下面是完整的封装类：
const ort = require('onnxruntime-node');
const sharp = require('sharp');
class Swin2SR {
async init(modelPath) {
this.session = await ort.InferenceSession.create(modelPath, {
executionMode: 'graph', graphOptimizationLevel: 'all'
});
}
async upscale(inputPath, outputPath) {
const { data, info } = await sharp(inputPath)
.removeAlpha()
.ensureAlpha(0)
.raw()
.toBuffer({ resolveWithObject: true });
// sharp输出的是RGBA，去掉alpha通道
const rgb = Buffer.alloc(info.width * info.height * 3);
let j = 0;
for (let i = 0; i < data.length; i += 4) {
rgb[j++] = data[i]; rgb[j++] = data[i + 1]; rgb[j++] = data[i + 2];
}
// HWC转CHW并归一化
const chw = new Float32Array(3 * info.width * info.height);
const plane = info.width * info.height;
for (let y = 0; y < info.height; y++) {
for (let x = 0; x < info.width; x++) {
const hwcIdx = (y * info.width + x) * 3;
chw[0 * plane + y * info.width + x] = rgb[hwcIdx] / 255;
chw[1 * plane + y * info.width + x] = rgb[hwcIdx + 1] / 255;
chw[2 * plane + y * info.width + x] = rgb[hwcIdx + 2] / 255;
}
}
const tensor = new ort.Tensor('float32', chw, [1, 3, info.height, info.width]);
const results = await this.session.run({ input: tensor });
// 输出CHW转回HWC写图，放大倍数为2
const out = results.output.data;
const scale = 2, outW = info.width * scale, outH = info.height * scale;
const outPlane = outW * outH;
const png = Buffer.alloc(outPlane * 3);
for (let y = 0; y < outH; y++) {
for (let x = 0; x < outW; x++) {
const idx = y * outW + x;
const hwcIdx = idx * 3;
const v0 = out[idx], v1 = outPlane + idx], v2 = out[2 * outPlane + idx];
png[hwcIdx] = Math.max(0, Math.min(255, v0 * 255));
png[hwcIdx + 1] = Math.max(0, Math.min(255, v1 * 255));
png[hwcIdx + 2] = Math.max(0, Math.min(255, v2 * 255));
}
}
await sharp(png, { raw: { width: outW, height: outH, channels: 3 } })
.png()
.toFile(outputPath);
return outputPath;
}
}
module.exports = Swin2SR;
调用方式很简单，实例化后先init加载模型，再传入图片路径即可：
const Swin2SR = require('./swin2sr');
(async () => {
const sr = new Swin2SR();
await sr.init('./models/swin2sr_x2.onnx');
await sr.upscale('input.jpg', 'output_x2.png');
console.log('超分完成');
})();
这里有个细节值得说明：输出张量的数值范围可能略微超出0到1，所以写回图片时要做clamp处理，否则画面会出现异常的纯黑或纯白噪点。另外Swin2SR是逐像素回归模型，输出是连续的浮点值，不要对结果做四舍五入以外的量化操作，会损失暗部细节。
大图分块处理与性能优化
直接把整张大图送进模型会遇到两个问题：一是显存或内存占用激增，一张2000x1500的图放大2倍，中间张量的内存占用轻松超过2GB；二是Transformer的注意力机制计算量与像素数呈平方关系增长，大图推理会慢得难以接受。实际工程中必须采用分块推理策略。
分块的思路是把原图切成互有重叠的小块，分别推理后再把重叠区域做融合，避免拼接处出现明显接缝。重叠宽度一般取16到32像素，融合时对重叠区域做线性加权过渡。示例逻辑如下：
async function tiledUpscale(sr, inputPath, outputPath, tileSize = 256, overlap = 16) {
const img = sharp(inputPath);
const { width, height } = await img.metadata();
const step = tileSize - overlap;
// 先对原图按tileSize分块，逐块调用推理
// 再用sharp的composite把结果块拼回大图
// 重叠区域按距离边缘的远近做线性权重融合
const cols = Math.ceil((width - overlap) / step);
const rows = Math.ceil((height - overlap) / step);
const composited = [];
for (let r = 0; r < rows; r++) {
for (let c = 0; c < cols; c++) {
const left = Math.min(c * step, Math.max(0, width - tileSize));
const top = Math.min(r * step, Math.max(0, height - tileSize));
const buf = await sharp(inputPath)
.extract({ left, top, width: tileSize, height: tileSize })
.png().toBuffer();
const tmp = `tile_${r}_${c}.png`;
await sr.upscale(buf, tmp);
composited.push({ input: tmp, left: left * 2, top: top * 2 });
}
}
await sharp({ create: { width: width * 2, height: height * 2, channels: 3, background: '#000' } })
.composite(composited)
.png().toFile(outputPath);
}
上面的简化版本没有做加权融合，直接用后处理的块覆盖前面的块，在平滑区域效果尚可，但纹理丰富的位置可能看出轻微的块边界。要彻底消除接缝，可以在composite之前对每块的四条边缘做alpha渐变蒙版，sharp生成蒙版的开销很小，值得加上。
性能方面有几个实测结论可以参考。CPU模式下（以8核桌面处理器为参照），128x128的小块单次推理大约300毫秒，256x256的块大约1.2秒，块越大单像素成本越低但内存压力越大，256是个不错的平衡点。如果需要GPU加速，Windows上可以改用onnxruntime-node的DirectML执行提供器：
const session = await ort.InferenceSession.create(modelPath, {
executionProviders: ['dml', 'cpu'] // DirectML优先，失败则回退CPU
});
需要注意的是，GPU加速对小块推理的提升有限，因为数据在CPU和GPU之间来回搬运的开销占比不小，分块尺寸提升到512之后GPU优势才会明显体现，通常能获得3到5倍的加速比。此外建议把会话实例做成全局单例，重复创建会话的开销接近每次几百毫秒，在Web服务场景下会显著拖慢响应速度。
常见问题排查
实际跑起来之后，有几个高频踩坑点需要提前了解。第一个是输出图像整体偏色或呈灰绿色，绝大多数情况是通道顺序搞反了，Swin2SR用的是RGB，而有些图像处理库默认BGR，两处转换逻辑要保持一致。第二个是onnxruntime-node在低版本Node上加载失败，建议使用Node 16以上版本，并确认安装过程中二进制下载完整，必要时删除node_modules重装。
第三个问题是内存峰值过高导致进程被杀，除了分块推理外，还可以在循环中及时释放中间Buffer，把Float32Array的生命周期控制在单个tile的推理周期内。JavaScript的GC回收大块内存有延迟，必要时可以显式置null并等待下一轮事件循环。最后一个容易被忽略的点是PNG输出体积，超分后的图像细节丰富，PNG压缩率会下降，如果用于网页展示，输出后追加一步WebP转换能把体积压缩一半以上：
await sharp('output_x2.png').webp({ quality: 90 }).toFile('output_x2.webp');
到这里，从模型转换到Node端推理、分块优化和问题排查的完整链路就打通了。这套方案的优势在于部署极简，一个Node进程加一个onnx模型文件就能提供超分服务，无论是嵌入Electron应用做本地画质增强，还是挂在图片处理服务后面做上传后处理，都非常合适。