目录
- 一、源码包准备
- 二、数据集准备
- 2.1 提取数据集名称
- 2.2 .txt报错问题
- 2.2.1 正确格式
- 2.2.2 错误格式
- 三、修改配置参数
- 四、训练及保存模型权重
- 4.1 训练
- 4.2 保存模型权重文件
- 五、模型推理测试
- 5.1 导入测试集
- 5.2 测试
- 5.3 测试结果
- 5.3.1 测试场景1
- 5.3.2 测试场景2
- 5.4 推理速度
- 5.4.1 CPU推理
- 5.4.2 GPU推理
- 六、总结
一、源码包准备
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参数生成噪声训练集才能达到效果最好,不同学者的训练测试结果有所不同是正常的。
后续有研究更深入的学者欢迎一起讨论学习。
总结不易,多多支持,谢谢!