git clone https://github.com/thstkdgus35/EDSR-PyTorch
以DIV2K数据集为例,我们将获得到的训练数据集按照以下文件放置自己选择好的文件目录中:
├── DIV2K
│ ├── DIV2K_train_HR
│ ├── DIV2K_train_LR_bicubic
│ │ ├── X4
│ │ │ ├── *****x4.png
│ ├── DIV2K_train_LR_bicubicL
│ │ ├── X4
│ │ │ ├── *****x4.png
为了减少代码修改工作,我们把对应的HR图像方法到DIV2K_train_HR文件夹中,对于赛道1我们把LR图像放到DIV2K_train_LR_bicubic
下的X4文件夹中。对于赛道2我们把LR图像放到DIV2K_train_LR_bicubicL
下的x4文件夹中。
对于测试数据我们按照以下方式创建并放入benchmark文件夹中就行:
├── DIV2K
├── benchmark
│ ├── test
│ │ ├── HR
│ │ ├── LR_bicubic
│ │ │ ├── X4
│ │ ├── LR_bicubicL
│ │ │ ├── X4
这里我们将自己写好的模型以小写英文命名放入EDSR项目中src文件夹下的model文件夹中中,例如如下ourmdel:
├── EDSR
│ ├── experiment
│ │ ├── src
│ │ ├── data
│ │ ├── loss
│ │ ├── model
│ │ │ ├── __init__.py
│ │ │ ├── edsr.py
│ │ │ ├── ourmodel.py
│ │ │ ├── ...
找到data中的__init__.py
文件, 我们打开后,我们修改如下几行:
for d in args.data_test:
if d in ['Set5', 'Set14', 'B100', 'Urban100', 'test']: #这里加入test数据集,要不然无法测试test中的数据
m = import_module('data.benchmark')
testset = getattr(m, 'Benchmark')(args, train=False, name=d)
我们找到div2k.py
和benchmark.py
的文件,打开后如下修改:
如果是进行赛道1则无需修改,直接下一步,如果是赛道2的话在这两个文件中的最后一行修改或加入
self.ext=('.png', 'jpg') #主要第一个.png是表示HR的文件格式,.jpg是表示LR的文件格式
针对赛道2我们修改一下srdata.py
中的代码:
self.input_large = (args.model == 'ourmodel') # 将ourmodel改成自己的模型名称
此处也可不改,在option.py文件中更改
# Model specifications
parser.add_argument('--model', default='rcan',
help='model name') # model下脚本名,选择为应用模型
我们找到src文件夹下的option.py
文件,打开修改以下参数:
parser.add_argument('--dir_data', type=str, default='../../../dataset', help='dataset directory') # 修改default后面的参数,修改成你放DIV2K文件的目录
parser.add_argument('--data_test', type=str, default='test', help='test dataset name') # 放入你设置的测试集名称如test
parser.add_argument('--epochs', type=int, default=300, help='number of epochs to train') # 修改你所需要的epoch大小
parser.add_argument('--batch_size', type=int, default=16, help='input batch size for training') # 修改你所需要的batch_size大小
parser.add_argument('--decay', type=str, default='200-400', help='learning rate decay type') #edsr的学习率策略是到了对应的轮数学习率减半,设置你要减半学习率的轮数,如果需要多个轮数降低就用-连接,例如 200-400代表200轮和400轮会分别降低一次学习率
训练就是打开demo.sh
文件把命令改成如下:
python main.py --model ourmodel --scale 4 --save ourmodel --pre_train ***/***/权重地址/ --chop --patch_size 192 #其中第一次训练可以不用--pre_train参数,chop参数是加速训练,可以不用
测试用以下命令在demo.sh
中:
python main.py --model ourmodel --scale 4 --data_test test --pre_train ***/***/权重地址/ --test_only # data_test 参数中的文件夹一定要在benchmark中有名字。 可以加--self_ensemble 参数生成多个结果来求和平均得到更高质量图像