平时大家做运动去模糊,第一反应就是端着 PyTorch 的预训练模型,连上服务器显卡去推理。这个思路没错,但如果你手上只有一台 macOS,不想额外租 GPU,又想快速验证 Restormer 在运动去模糊上的效果,就会发现常规路线其实有点绕:环境里要装一堆依赖,MPS 后端对某些算子支持又不够稳定,跑起来总有一种“能用但不好用”的别扭感。我后来把整套流程切到 Swift + MLX 上,用 mlx-restormer-swift 在 macOS 本地完成运动去模糊推理,整个过程反而顺畅很多,尤其 fp32 版本在精度和调试体验上特别让人安心。
这篇文章就把我搭建这套推理流程的完整思路和实操过程拆开讲清楚。内容包括为什么选 Restormer 而不是传统 CNN 结构、MLX 在 macOS 上的优势、fp32 与半精度怎么取舍、模型权重怎么转换、Swift 推理主流程怎么写,以及我实际踩过的坑。适合想在 macOS 上用自研或开源模型做图像恢复推理的开发者,也适合被 Python 环境折腾烦、想试试 Swift 方案的人。
1. 内容整体设计与思路拆解
1.1 为什么选 Restormer
运动去模糊本质上是图像恢复任务,输入是带有运动模糊的退化和模糊图像,输出是清晰的潜在图像。传统方法里,CNN 结构比如 U-Net 或者残差网络也能做,但缺点是感受野有限。模糊轨迹往往覆盖很大范围,一帧图像里远处和近处的运动关系是长距离依赖,CNN 靠堆叠卷积层来扩大感受野,计算量和效果都不够理想。
Restormer 的核心优势就是引入 Transformer 结构,但并不是像 ViT 那样把图像切 patch 再做空间注意力。它在通道维度上计算多头注意力,同时能捕捉跨通道的全局依赖,计算复杂度又不会随着图像分辨率的平方增长。这个特性对高分辨率图像的恢复特别关键。
在做运动去模糊时,图像的边缘纹理和模糊轨迹之间是强关联的,Restormer 在多个尺度上交互信息,可以让模型更“聪明”地判断哪些区域该锐化、哪些区域该抑制伪影。这一点在真实场景里很实用:比如手机抓拍运动物体、运动相机画面抖动,模糊模式往往不是均匀的,Restormer 处理这种空间变化的模糊比普通 CNN 稳定得多。
1.2 为什么在 macOS 上用 MLX 而不是 PyTorch
PyTorch 其实早就支持 MPS 后端,理论上也能在 macOS 上做 GPU 推理。但实际体验有几个问题:
第一,环境管理。为了跑一个推理脚本,往往要装 Python、PyTorch、NumPy、OpenCV,虚拟环境一不小心就崩。第二,MPS 性能。虽然 MPS 对主流算子支持得不错,但遇到自定义算子或者某些归一化层,经常回退到 CPU,速度优势就没了。第三,部署形态。Swift 应用里想调用 Python 模型,要么走 Socket 传入传出,要么跑子进程,这层胶水代码很别扭。
MLX 就不太一样。它是构建在 macOS 统一内存架构上的机器学习框架,数组对象MLXArray和 NumPy 的接口很像,上手成本低,同时原生支持 Swift。用 Swift 直接加载权重、建计算图、跑推理,不需要桥接层,整个推理程序可以很干净地集成到命令行工具或桌面应用里。
统一内存意味着 CPU 和 GPU 共享同一块内存,数据不用拷贝来拷贝去。这个设计在处理图像输入输出的场景下特别香,前处理和后处理可以直接操作同一块缓冲区。
1.3 fp32 版本的意义
这个项目标题里明确带 fp32,说明权重和计算精度都固定在 32 位浮点。很多人可能会问:现在模型压缩这么流行,fp16 不是更快更省内存吗?为什么还要专门做 fp32?
考虑两个原因。第一是精度可控。fp32 在数值表示上没有额外误差,训练时怎么算,推理时就怎么算,调试时不容易出现“模型输出发绿”“结果有噪点”这种精度引起的诡异问题。第二是通用性。MLX 虽然支持 fp16,但某些自定义层在 fp16 下会有数值溢出风险,尤其是通道注意力里的 softmax、层归一化。对于开发者来说,先把 fp32 跑通,再考虑混合精度优化,这是更稳的推进方式。
代价也很明显:fp32 权重体积是 fp16 的两倍,内存占用也更高。但只要不是大规模并发推理,单张图像在 macOS 本地上跑,内存压力其实是可控的,后面我会给出具体估算。
1.4 整体实现思路
我把整体实现拆成四个环节:
- 模型准备:先拿到预训练权重,搞清楚网络结构,确认输入输出张量的形状和归一化方式。
- 环境搭建:建一个 Swift Package,导入 MLX 依赖。
- 推理管线:加载权重到模型,写图像前处理、模型 forward、后处理、结果保存。
- 封装测试:把推理逻辑封装成命令行工具,测单张图片,再测批量目录。
整个流程不依赖显卡,不需要 CUDA,单靠 M1/M2/M3 芯片上的 Metal GPU 就能跑。下面依次拆解每个环节的核心细节。
2. 核心细节解析与实操要点
2.1 Restormer 的关键模块与权重映射
Restormer 整体是一个类似 U-Net 的结构,有编码器、瓶颈、解码器,不同层之间用跳跃连接。每个 Transformer Block 里,最核心的模块是多头通道注意力,英文缩写是 MDTA。它把层归一化后的特征图在通道维度分组,每组内部通过矩阵乘法计算注意力图,再用注意力图调整原始特征。
这个设计的厉害之处在于:空间维度保持不变,避免了传统空间自注意力随着图像尺寸变大而带来的爆炸性计算;但跨通道的信息交互依然保留得很完整。比如一张运动去模糊的输入图,红色通道里看到的模糊纹理可能和蓝色通道里看到的纹理高度相关,MDTA 能把这种跨通道线索利用起来。
Gate mechanism 也值得一提。Restormer 的 Feed-Forward 部分引入了门控线性单元,类似 Gated Linear Unit,输出是两条分支相乘。这让模型在传播信息时多了一个“阀门”,对去模糊任务中保留结构边缘很有帮助。
权重映射时,需要注意卷积核的排列顺序。PyTorch 的卷积权重形状是[out_channels, in_channels, kh, kw],MLX 的卷积如果采用Conv2d模块,权重的排列也一样,所以大部分情况下可以直接搬运。唯一要小心的是 LayerNorm 里的weight和bias,形状是[num_channels],这两个大小写必须和模型定义完全对齐。
2.2 图像预处理与后处理规格
做运动去模糊时,模型训练阶段通常会把像素除以 255 归一化到[0,1],然后减均值。不同仓库的均值设定不一样,有的直接用0.5减,有的用 ImageNet 统计。这里不能想当然,我之前就吃过亏,下面细说。推理阶段必须还原训练时的预处理方式,否则结果会明显偏色。
输入尺寸也有讲究。Restormer 内部有多层下采样,最深的特征图分辨率会是输入的1/8或者1/16。为了让下采样边界整齐,输入图像的宽高最好能被 16 整除。如果原图尺寸不满足,我前处理时会在边缘做 padding,或者把图像切割成重叠块,推理完再将输出拼接回原尺寸。
后处理相对简单:输出张量先恢复到[H, W]区间,再加均值、乘 255,最后截断到[0,255]。截断这步很多人会漏掉,网络输出可能会出现边缘溢出,尤其经过多尺度上采样之后,某些像素在 1 附近波动,不截断直接保存 PNG,画面上就会出现奇怪的亮点。
2.3 Swift 中操作 MLX 张量的几个习惯
MLX 里的核心类型是MLXArray,操作风格和 NumPy 非常像。你可以在 Swift 里写x + y、x.reshaped()、x.transposed(),不会有 Swift 标准库类型转换的割裂感。上手时我用了一个很简单的心智模型:写 Python 时候怎么操作 NumPy,写 Swift 就怎么操作 MLXArray。
有几个细节值得注意:
- 张量维度顺序。PyTorch 和 MLX 都遵循
CHW的批处理格式,即[B, C, H, W]。图像从HWC转到CHW时,需要用transposed。 - 不要写逐像素循环。MLX 底层虽然做了图优化,但 Swift 层如果手动遍历图像每个像素,性能会差很多。能用
divide、subtract、clipped的整体算子就不要用循环。 - 归一化可以用
asType先统一数据类型。输入图像如果是UInt8,得先转成Float32,再除以 255.0。
2.4 fp32 内存估算与 batch 设计
fp32 的内存占用是可以直接算出来的。一个像素有 3 个通道,每个通道 4 字节。如果输入分辨率是[1, 3, 512, 512],那输入本身占 3MB。假设置于编码器最低分辨率处的特征图通道数为 512,空间尺寸为[1, 512, 64, 64],就有 8MB。Transformer 内部的多头注意力还会产生中间张量,比如注意力矩阵,如果通道分组为 8,每个头内部矩阵尺寸可能是[1, 512, 64, 64]这样的量级。
所以保守估计,一张 512×512 的输入图,fp32 推理的峰值内存可能在 500MB 到 1GB 之间。这个数值对 Mac 的统一内存来说并不夸张,但如果你要批量处理上百张图,就得考虑一次只送一张,或者控制并行数。batch size 设成 1 最稳,因为去模糊任务通常不依赖 batch 内的统计信息,MLX 也更容易做内存复用。
3. 实操过程与核心环节实现
3.1 准备环境
这里假设你已经有了 Xcode 或者 Command Line Tools。先确认基础工具:
xcode-select --install swift --version能在终端看到 Swift 版本输出就行。接着新建一个 Swift Package:
mkdir RestormerDemo cd RestormerDemo swift package init --type executable然后用 Swift Package Manager 添加 MLX 的 Swift 依赖。这一步不用手动下载源码,Package 解析时会自动拉取。我测试时用的是 0.16.0 以上版本,更早的版本接口有细微差异,建议直接使用较新的稳定版。在Package.swift里的声明类似这样:
// swift-tools-version:5.9 import PackageDescription let package = Package( name: "RestormerDemo", platforms: [.macOS(.v13)], dependencies: [ .package(name: "MLX", path: "path/to/mlx-swift") ], targets: [ .executableTarget( name: "RestormerDemo", dependencies: [.product(name: "MLX", package: "MLX")] ) ] )如果是从官方仓库克隆的 MLX Swift 项目,直接用path引用本地目录编译速度更快。每次改完依赖记得跑一次。
swift build能编译通过,说明基础环境没问题。
3.2 模型权重转换的两种方式
拿到 PyTorch 预训练权重后,需要把它转换成 MLX 能加载的格式。我试过两条路线,都走得通。
方式一:先把 PyTorch 权重导出为独立的.npy文件。
import torch import numpy as np checkpoint = torch.load("restormer_motion_deblur.pth", map_location="cpu") state_dict = checkpoint["state_dict"] if "state_dict" in checkpoint else checkpoint import os os.makedirs("converted_weights", exist_ok=True) for key, value in state_dict.items(): arr = value.detach().float().numpy() # 文件名把点换成下划线,方便 Swift 里按字符串索引 filename = key.replace(".", "_") + ".npy" np.save(os.path.join("converted_weights", filename), arr) print("导出完成,共", len(state_dict), "个权重")方式二:如果用 Safetensors,加载更安全,也省去担心 pickle 的安全问题。转换思路一样,遍历 key,把每个张量取出来保存成独立文件或者打包成一个 zip。
在 Swift 端,加载权重时用MLXArray从.npy文件读取。每个 key 的值会变成一个MLXArray,再塞进模型参数字典。关键点在于模型结构里所有层的名字必须和导出的 key 一一对应。如果 Swift 重构模型时层名不一致,权重加载就会错位。
为了方便调试,我通常会把模型打印出来和 PyTorch 的网络结构对比一眼:
for (key, value) in parameters { print(key, value.shape) }确认每个名字和形状都对得上,再继续后续操作。
3.3 Swift 推理主流程代码
下面展示一个最小可用的 Swift 推理流程。这里把模型类定义简化成Restormer,你实际使用时只要保证它的callAsFunction接收MLXArray,输出也是MLXArray即可。
import MLX import MLXNN import Foundation // 读取模型权重,返回 [字符串: MLXArray] func loadWeights(from directory: String) throws -> [String: MLXArray] { var dict: [String: MLXArray] = [:] let fileManager = FileManager.default let files = try fileManager.contentsOfDirectory(atPath: directory) for file in files where file.hasSuffix(".npy") { let key = file.replacingOccurrences(of: "_", with: ".") .replacingOccurrences(of: ".npy", with: "") let path = directory + "/" + file let array = try MLXArray(contentsOf: URL(fileURLWithPath: path)) dict[key] = array } return dict } // 最简单的图像预处理 func preprocess(_ image: MLXArray) -> MLXArray { // 输入形态:HWC,转成 CHW 并归一化 let chw = image.transposed(axes: [2, 0, 1]) let floatImg = chw.asType(.float32) let normalized = floatImg / 255.0 return normalized[0, 0..<3, 0..<64, 0..<64] // 这里用固定尺寸做演示,实际按模型输入裁剪 } // 后处理:从 CHW 回到 HWC,并截断 func postprocess(_ output: MLXArray) -> MLXArray { let clipped = output.clipped(lower: 0.0, upper: 1.0) let hwc = clipped.transposed(axes: [1, 2, 0]) let uint = (hwc * 255.0).asType(.uint8) return uint } let model = Restormer() let weights = try loadWeights(from: "converted_weights") model.update(parameters: weights) let inputImage = try MLXArray(contentsOf: URL(fileURLWithPath: "input.png")) let inputTensor = preprocess(inputImage) let outputTensor = model(inputTensor) let outputImage = postprocess(outputTensor)代码里我对输入尺寸做了 hardcode,这是为了让流程先跑通。实际项目中,图像要经过尺寸对齐、裁剪或 padding,才能送进模型。
Swift 的MLXArray可以直接从图片文件初始化吗?我的做法是先用 Core Image 或 ImageIO 把 PNG 解析成像素 Buffer,再转成MLXArray。如果你不想引入比较重的依赖,也可以用我自己封装的一个最小读取函数,先拿到 RGBA 数据,再扔掉 Alpha 通道,组装成 MLXArray。
3.4 封装成命令行工具并集成测试
最小推理流程能跑通后,我建议马上把它封装成命令行工具,这样后续批量测试会轻松很多。比如定义一个参数结构:
RestormerDemo --input blur01.png --output sharp01.png --weights converted_weights在main.swift里用CommandLine.arguments手动解析,不需要额外引库。核心逻辑就是读取参数、初始化模型、循环处理文件、清理输出。
批量测试时,我直接遍历一个目录里的所有图片,一张一张推理,最后统计成功率。这里的“成功率”不只是有没有跑完,还包括输出图像的像素是否正常,比如有没有全黑、全白、明显伪影。
集成测试阶段,我会准备三张不同难度的图片:
- 一张简单场景,比如静止被摄物但相机轻微抖动;
- 一张中等场景,画面里有平移的汽车;
- 一张复杂场景,包含树木和行人,运动轨迹不规则。
简单场景要求输出几乎无感,中等场景要求纹理清晰,复杂场景只要整体没有明显假色和重影,就算通过。
4. 常见问题与排查技巧实录
4.1 常见问题速查表
| 现象 | 可能原因 | 排查思路 |
|---|---|---|
| 输出全黑 | 后处理没有截断或归一化错误 | 打印输出张量的 min/max,确认是否在合理范围 |
| 输出全白 | 数值溢出或均值加错 | 检查预处理是否重复减均值,确认 float 精度 |
| 颜色偏色严重 | 通道顺序错误 | 检查 HWC 转 CHW 的 transpose 顺序,RGB 是否变成 BGR |
| 输出尺寸不对 | 模型输入和实际输入不匹配 | 打印模型各层输出形状,逐一对齐 |
| 内存暴涨甚至崩溃 | fp32 峰值内存过高 | 降低输入尺寸,关闭后台 app,尝试流式分块推理 |
| GPU 参与度低 | 部分算子在 Metal GPU 不支持 | 用 CPU 模式跑一遍看速度和结果差异 |
| 边缘有黑边 | padding 没有还原 | 记录 padding 大小,后处理裁剪回来 |
| 权重加载失败 | key 名称不匹配 | 对比 PyTorch state_dict 和 Swift 模型层名 |
4.2 踩坑记录与排查思路
先说一个我印象最深的坑:卷积权重维度相同,但模型输出全是噪点。排查很久发现,问题出在 LayerNorm 的 epsilon 上。PyTorch 的 LayerNorm 默认eps=1e-5,MLX 的一些实现里如果用默认值也是一样,但如果你在加载权重前把模型参数重新初始化过,某些初始化函数会把 epsilon 改成1e-6。训练和推理的 epsilon 不一致,结果就会差很微妙,肉眼看起来就是颜色发花。排查方法是直接打印中间层输出,定位到第一个出现异常的特征图,然后对比代码里 LayerNorm 的参数。
第二个坑是尺寸对齐。训练时 Restormer 为了效率,一般会把输入缩放到 512×512 或者 256×256。但实际用的时候,摄影原图可能是 4000×3000。直接 reszie 到 512 会损失大量细节,去模糊效果会打折扣。我后来采用滑窗方式:把大图切成 512×512 的重叠块,块与块之间重叠 64 像素,推理完再按重叠区域做淡入淡出融合。这样既能保留细节,又不会让 GPU 内存爆掉。融合系数我用了线性加权,重叠中心权重大,边缘权重小,效果挺好。
第三个坑是 CPU 和 GPU 的行为差异。同一份权重,在 CPU 上和 GPU 上的结果会有细微差别,差的不是很大,但如果你有精确比对需求,就要固定设备来跑。MLX 里默认会自动调度,但你可以通过设置让整个图跑在 CPU 上。这样做虽然慢一些,但结果更可复现。
第四个坑是关于内存的。fp32 模型虽然精度高,但加载后如果同时保留 PyTorch 转换出来的.npy文件、Swift 里的模型参数、中间特征图,内存占用会急涨。我的经验是权重转成.npy后就可以把原始.pth文件移到归档目录,Swift 进程里只保留一份模型参数,用完一个中间变量就尽早置空。
最后一个小技巧:如果模型推理速度很慢,不一定是 MLX 的问题,可能是你的输入尺寸太大导致下采样后特征图依旧很大。比如 1024×1024 输入,编码器最低层如果还保留 256×256 的空间尺寸,注意力计算就会很重。这时候与其换模型,不如先把输入调整到 512 左右,观察峰值内存和速度是否同时下降,再判断是否值得继续优化。
实际用下来,我发现这套 Swift 方案最舒服的地方不是跑分多高,而是调试链路短。Python 那边出问题了要同时查模型、查 NumPy、查环境,Swift 这边一旦编译通过,剩下的问题几乎全集中在张量形状和数值范围上。你可以在swift里像调试普通程序一样打断点,直接看中间结果,最后把推理封装成命令行工具,也能很方便地接入其他工作流。
如果你也正准备在 macOS 上做去模糊或者类似图像恢复的推理,我建议先把流程跑起来,不要纠结精度和速度。第一版只要能快速验证“模型能加载、输入能进去、输出能存出来”,再逐步加入分块、批处理这些优化。用 MLX 的好处就是每一步都足够透明,适合像我这样喜欢把整个技术栈都抓在自己手里的人慢慢折腾。