图像去噪——PMRID训练自己数据集及推理测试（详细图文教程）

一、源码包准备

RMRID最原始的版本是MegEngine框架的，后有了Pytorch版本，本教程主要是Pytorch版本的。

官网MegEngine版本源码：PMRID(MegEngine)

官网Pytorch版本源码：PMRID(Pytorch)

论文地址：论文

我自己提供的源码包是在官网基础上修改过的，其中有一些我添加的代码，我提供的源码包链接为：网盘源码包 ，提取码：he6m

我提供的源码包下载并解压后的文件内容如下：

二、数据集准备

数据集有很多，看自己需要哪个数据集，官网提到的Kaggle数据集下载链接为：Kaggle

打开网址后直接下载即可，如下：

如果你想用.RAW格式的原始数据训练，在读入数据这一部分，得修改代码，读入.RAW格式数据的代码如下，关于读取.RAW格式数据和.dng格式数据的方法，查看另外一篇博文读取.RAW数据和.dng数据

如果数据集是.jpg，.png，.bmp等格式的8bit数据，则保持默认即可，如下：

2.1 提取数据集名称

在提供的脚本文件中，通过generate_list_sidd.py，将含噪图像名称和真实无噪图像名称提取到一个.txt文件中，如下：

从生成好的train_list_sidd.txt文件中随机剪切一部分到其它两个文件Sony_val_list.txt和Sony_test_list.txt中，作为验证集和测试集，如下：

2.2 .txt报错问题

上一步骤中提取图片名称到.txt文件中，注意.txt文件中最后一行的空行一定要删除，不然就会报错，如下：

2.2.1 正确格式

2.2.2 错误格式

三、修改配置参数

可以修改的超参数如下

训练轮数：epoch
训练批次大小：batch_size
学习率：learning_rate
使用GPU或CPU：device
日志文件路径：logs_path
模型保存路径：params_path
训练集路径：train_list_path
验证集路径：value_list_path
是否加载预训练权重文件：is_load_pretrained
预训练训中路径：pretrained_path

源码包中具体修改的地方如下，在mian.py脚本中：

四、训练及保存模型权重

4.1 训练

修改好上面参数后就可以开始训练了，训练前先打开训练命令，如下：

运行脚本main.py后就开始训练了，如下：

4.2 保存模型权重文件

训练过程中会将每个epoch的模型权重文件保存在params路径下，如下:

模型大小为3.98M

五、模型推理测试

训练好后，进行推理前先打开测试命令，如下：

5.1 导入测试集

由于代码架构设计，验证集的对应的变量名和测试集对应的变量名是一个，所以直接将测试集路径替换掉验证集路径，如下：

5.2 测试

修改好上面参数后，运行脚本main.py，会在output路径中生成测试结果：

5.3 测试结果

5.3.1 测试场景1

左侧为原始含噪图像，右侧为模型测试结果

5.3.2 测试场景2

5.4 推理速度

5.4.1 CPU推理

我测试的电脑为i7-12700H 2.3GHz，测试图片大小为256*256，CPU上测试的平均推理时间为40ms，部分如下：

5.4.2 GPU推理

GPU上的平均推理时间为：6ms，部分如下：

六、总结

以上就是使用PMRID训练自己数据集及推理测试的详细图文教程，由于推理速度达不到要求，后续我没有再继续研究，且本方法要结合相机的iso参数生成噪声训练集才能达到效果最好，不同学者的训练测试结果有所不同是正常的。

后续有研究更深入的学者欢迎一起讨论学习。

总结不易，多多支持，谢谢！