NAFNet
NAFNet 是一个用于图像恢复的卷积网络,它把典型 UNet 模块里的非线性激活函数去掉,换成逐元素相乘。LibreYOLO 支持它的一个任务,restoration,并发布了一个在 SIDD 上训练的真实图像去噪检查点。
- 任务
- restore
- 尺寸
- s, l at 256 px
- 安装
pip install libreyolo- 支持层级
- 已支持,自 v 起。起支撑作用的可训练家族:在 CI 里保持绿色,功能视情况落地。
- 许可
- 代码采用 MIT,权重采用 MIT。商用
安装
NAFNet 不需要任何可选 extra。它导入的一切都在基础安装里。
pip install libreyolo预测
权重在首次使用时从 Hugging Face 下载,并缓存在本地。
from libreyolo import LibreYOLO model = LibreYOLO("LibreNAFNetl-restore-sidd.pt")result = model("noisy.jpg", save=True) restored = result.restoredprint(restored.array.shape)libreyolo predict model=LibreNAFNetl-restore-sidd.pt source=noisy.jpg save=Truefrom libreyolo import LibreYOLO model = LibreYOLO("LibreNAFNetl-restore-sidd.pt")result = model.predict("noisy.jpg") result.restored.save("denoised.png")返回的 Results 对象为这个家族只带一个字段 restored,它是原始画布上一张稠密的
HWC uint8 RGB 图像;没有检测框可以遍历。save=True 会把这张恢复后的图像直接写入
磁盘,而不是在输入图上画一层标注。conf、iou 和 max_det 为了和其他每个家族
保持签名一致而被接受,但不起作用,因为图像恢复不产生任何可供过滤的检测结果。数据
源、流式处理和结果处理见预测。
变体
两种宽度共享这套架构:s(宽度 32)和 l(宽度 64),都围绕 256 px 的训练 patch
构建。无论尺寸如何,预测和验证都在原生图像分辨率下运行,只把图像补齐到网络的下采
样因子。目前只发布了 l 这个宽度,它是一个在 SIDD 上训练的真实图像去噪检查点
(checkpoint)。
训练
NAFNet 在你自己的成对退化/干净图像上微调:一个数据集 YAML,指向存放退化图像的
inputs/<split>/ 文件夹和存放干净目标的 targets/<split>/ 文件夹,两者按文件名主
干匹配。degradation 和 dataset 是可选字符串,记录在保存的检查点上用于溯源;它
们不参与训练。
from libreyolo import LibreYOLO model = LibreYOLO("LibreNAFNetl-restore-sidd.pt")model.train(data="my-dataset.yaml", epochs=100, imgsz=256, batch=16, lr0=1e-3)libreyolo train model=LibreNAFNetl-restore-sidd.pt data=my-dataset.yaml \ epochs=100 imgsz=256 batch=16 lr0=1e-3from libreyolo import LibreYOLO # degradation 和 dataset 会记录在保存的检查点上,# 它们不改变训练的内容model = LibreYOLO("LibreNAFNetl-restore-sidd.pt")model.train( data="my-dataset.yaml", epochs=100, degradation="denoise", dataset="MyDataset",)libreyolo train model=LibreNAFNetl-restore-sidd.pt data=my-dataset.yaml \ epochs=100 device=0,1 batch=32不去动它时,训练器会跑 100 轮,用 AdamW、lr0=1e-3、批大小 16、256 px 裁剪,并在
PSNR 连续 50 轮没有提升后早停。这个家族没有 LoRA 路径:lora=True 会直接报错而不
是运行,因为 NAFNetTrainer 从不启用适配器微调。
训练期间网络跑的是普通的全局平均池化。NAFNet 那个只用于推理的窗口化局部池化 (Test-time Local Converter)会在第一轮开始前被摘掉,训练结束后再装回去,因为在固 定窗口的局部池化上做反向传播,和这个检查点在推理时的用法对不上。
数据集、数据增强、多卡训练和日志记录器见训练。
验证
val() 返回一个字典,里面有 metrics/PSNR 和 metrics/SSIM,在 RGB 上、按整张有
效画布计算:SSIM 用的是 11x11 的高斯窗口,sigma 为 1.5,而用于挑选最佳检查点的
fitness 就是 PSNR 值。data 指向的是训练时用的同一种成对图像数据集格式。
from libreyolo import LibreYOLO model = LibreYOLO("LibreNAFNetl-restore-sidd.pt") # val() 返回的是普通 dict,不是对象metrics = model.val(data="my-dataset.yaml") print(metrics["metrics/PSNR"])print(metrics["metrics/SSIM"])libreyolo val model=LibreNAFNetl-restore-sidd.pt data=my-dataset.yaml导出
| 任务 | ONNX | TorchScript | ExecuTorch | TensorRT | OpenVINO | Paddle | MNN | RKNN | ncnn | TFLite | CoreML | Core AI |
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| restore | restore to ONNX:支持 | restore to TorchScript:支持 | restore to ExecuTorch:支持 | restore to TensorRT:支持 | restore to OpenVINO:支持 | restore to Paddle:不支持 | restore to MNN:不支持 | restore to RKNN:不支持 | restore to ncnn:支持 | restore to TFLite:不支持 | restore to CoreML:不支持 | restore to Core AI:支持 |
导出的产物按文件后缀经由 LibreYOLO() 重新加载,所以一个 .onnx 或 .engine 文
件的表现和检查点一样,返回同样的 Results,由 restored 承载输出图像。NAFNet 以
固定的空间分辨率导出:imgsz 必须能被网络的下采样因子整除(两种架构宽度都是
16),而且 dynamic=True 时只有 batch 这一维是动态的;高和宽在导出时就固定下来。
from libreyolo import LibreYOLO model = LibreYOLO("LibreNAFNetl-restore-sidd.pt")model.export(format="onnx", imgsz=256)model.export(format="tensorrt", imgsz=256, half=True)libreyolo export model=LibreNAFNetl-restore-sidd.pt format=onnx imgsz=256libreyolo export model=LibreNAFNetl-restore-sidd.pt format=tensorrt imgsz=256 half=Truefrom libreyolo import LibreYOLO # 工厂按文件后缀分发,所以导出的产物加载方式和任何检查点一样,# 返回的也是同一个 Results 对象model = LibreYOLO("LibreNAFNetl-restore-sidd.onnx")result = model("noisy.jpg") result.restored.save("denoised.png")检查点
这个家族已发布的全部权重文件。
| 文件 | 输入(px) | 权重许可 |
|---|---|---|
| restore | ||
| LibreNAFNetl-restore-sidd.pt | mit | |
上面的每个文件目前都在 LibreYOLO 组织中,并会在首次使用时下载。
许可证
请检查你所下载的具体权重在 Hugging Face 仓库中的许可。LibreYOLO 组织里的每个检查点都附有许可,同一家族内也不一定相同。该仓库是权威来源;以下摘要说明本页上次验证时适用的情况。
这里只说明涉及的许可证,不构成法律意见。如果答案对商用很重要,请自行阅读许可证并咨询法律顾问。
- 原始工作
- NAFNet, Megvii
- 上游许可
- MIT
- LibreYOLO 代码
- MIT
- 权重
- 采用 MIT 许可,重新发布在 huggingface.co/LibreYOLO
- 解读
- MIT is a permissive license, so these weights can be used in commercial and closed-source products. It asks you to keep the copyright notice and license text with any copy you redistribute, and places no other obligation on your own application code. Part of the training pipeline is ported from BasicSR under Apache-2.0, which additionally grants a patent license. The published checkpoint is trained on the Smartphone Image Denoising Dataset (SIDD), itself MIT-licensed.
引用
@article{chen2022simple,
title={Simple Baselines for Image Restoration},
author={Chen, Liangyu and Chu, Xiaojie and Zhang, Xiangyu and Sun, Jian},
journal={arXiv preprint arXiv:2204.04676},
year={2022}
}复制自作者在 github.com/megvii-research/NAFNet#citations 上提供的引用块。