旺丁旺财杂货铺
生活/旅行/效率/电器/装修
🎯 Intro 前记
两年前,由于要将扫描的渣画质PDF书本变清晰,铺主折腾了AI画质增强/超分辨率好一段日子。这两年间,软件由upscayl换到realesrgan-gui,再到加装了chaiNNer;模型由Tencent NCNN到PyTorch。一路体验,发现Tencent NCNN实在强,连嵌入式的小设备都能运行得好好的,不像PyTorch离开CUDA都难。于是,铺主就想将Real-ESRGAN的pth模型转成NCNN到realesrgan-gui用。没想到铺主踩了近些年最大的坑,绕了近年来最大的弯,踏进中文互联网好大的一块空白。幸好,最近总算爬上来了。这次就聊聊怎么将Real-ESRGAN的pth模型转成realesrgan-ncnn-vulkan可用NCNN吧
🐛 PyTorch到NCNN的坑
虽然AI很火,超分辨率又是AI图像生成技术的底座之一,但研究的人实在是真TM少。而github.com/Phhofm的Philip Hofmann算是全球独一档的大神,以一自之力为全球贡献了不少的超分辨率模型。其中我用得最多的是他的4xNomosWebPhoto_esrgan,它是为数不多 超分之后中文字体不产生伪影(不同笔画之间生成连线)的模型,堪称扫描PDF图像增加真神。
目前openmodeldb.info和huggingface仅提供了4xNomosWebPhoto_esrgan模型pth格式的下载。Philip Hofmann在自己的github.com/Phhofm/models/releases/tag/4xNomosWebPhoto_esrgan里提供了包含NCNN在内多种格式的模型。可在作者那下载NCNN格式之后,realesrgan-ncnn-vulkan载入模型便报错。我还记当时大费周章之后得到的竟然是这结果,差点
为此,铺主曾经试用chaiNNer里PyTorch转NCNN的节点、PyTorch转ONNX的节点,再将ONNX交到onnx2ncnn.exe转换。然而,无一不以翻车告终。原本铺主放弃了一段时间,直至最至最近,鬼使神差之下,与AI Chat和过程得知可以利用PNNY将PyTorch转成NCNN。然而中文PNNY的使用,介绍简单得像食饭屙屎一般轻易,实际上一条有用的条目也没有,按介绍的来只有报错,没有结果。幸好,PNNY是python的库,铺主打开Vscode,花了半天,久经波折之下终于获得realesrgan-ncnn-vulkan可用的NCNN,下面就讲讲思路吧。
🚀 Python利用PNNY将模型由PyTorch转成NCNN
1. 配置与初始化
- 路径设置:获取脚本所在目录 [SCRIPT_DIR],定义输入模型文件名[INPUT_PTH_FILENAME] 和输出前缀 [OUTPUT_PREFIX_FILENAME]。
- 参数设定:设置 dummy input 的形状 [INPUT_SHAPE] (1, 3, 64, 64) 以及是否使用 GPU [USE_GPU]。
上面[INPUT_SHAPE] (1, 3, 64, 64) 的参数是通用模型realesrgan-4xplus的参数,如果训练模型的参数不一样是需要修改的,可以找一下模型github的出处,看YML文件有没有相关配置,或者写另外代码验证。
2. 模型定义 (替代 Basicsr)
直接定义了 RRDBNet 及其子模块,避免依赖庞大的 basicsr 库:
- [ResidualDenseBlock]: 定义残差密集块,包含5个卷积层和 LeakyReLU 激活函数,并使用 [default_init_weights]进行权重初始化。
- [RRDB]: 由3个 [ResidualDenseBlock] 组成的残差块。
- [RRDBNet]: 主网络架构,包含初始卷积、多个 RRDB 主体块、上采样模块(两次近邻插值+卷积)和最终输出卷积。
3. 主执行流程 ([main]函数)
A. 环境准备
- 切换工作目录 保存当前工作目录,并将工作目录切换到脚本所在目录 [SCRIPT_DIR]。这是为了确保 PNNX生成的中间文件路径干净,不包含绝对路径分隔符。
B. 模型加载 ([load_model](javascript:void(0)))
- 设备选择:根据 [USE_GPU]和 CUDA可用性决定使用 cuda 还是 cpu。
- 加载权重: 使用 torch.load 加载 .pth文件。 智能提取状态字典:优先查找 params_ema,其次是 params,最后是整个 checkpoint。 键名清洗:移除权重键名中的 module. 前缀以匹配模型定义。 加载状态字典并处理缺失或意外的键警告。 将模型移动到指定设备并设置为评估模式 (eval())。
C. PNNX 转换
- 生成 Dummy Input:创建一个符合 [INPUT_SHAPE]的随机张量作为追踪输入。
- 调用 PNNX: 尝试导入 pnnx 模块。 调用 pnnx.export() 将 PyTorch 模型转换为中间表示。 容错处理:首先尝试标准调用方式,如果因参数类型错误失败,则尝试备用调用方式(将 dummy_input 包裹在列表中)。
D. 后处理与优化
如果转换成功:
- 修复 Param 文件 ([fix_param_file]: 读取生成的 .ncnn.param 文件。 将输入节点名称 in0 替换为 data。 将输出节点名称 out0 替换为 output。 目的:使生成的模型兼容 realesrgan-ncnn-vulkan 推理工具的默认输入输出命名规范。
- 清理中间文件 ([cleanup_intermediate_files]: 保留最终需要的 .ncnn.param 和 .ncnn.bin 文件。 删除 PNNX 生成的中间文件,如 .onnx, .pt, .pnnx.py, .pnnx.bin 等,保持目录整洁。
E. 收尾
- 恢复工作目录:无论成功与否,最后都将工作目录恢复为原始目录。
- 输出结果:打印转换状态、修复信息和最终生成的文件路径。
由于代码太长,我便不贴在这里了。不过按上面的提示,AI都能生成不错的代码了,哈哈。
✏️WQ: 后记
有位旧友总喜欢说“凡走过,必留痕迹”。我也想不到,当初只是想搓一个输入PDF超分计算后输出AI增强变清晰PDF的小软件,让我绕了这么多路,累积了这么多奇奇怪怪的经验。我不知道我这个淋过雨的人写的这篇分享会不会让为其它人撑起一把伞。相关的 实例代码已经上传到Gitee(地址见文末),随自取缘。好了,这次就写到这了,若然觉得这篇文章能帮到你,希望大家点赞和评论留言,这能帮到我不少,感谢感谢!

🌍https://gitee.com/ginkdu/py-torch2-ncnn_convert_with_pnnx