hr_size = 256
target_net = DnCnn(layer_num=20)
target_net.load_state_dict(checkpoint)
model = SuperviseModel({
'net': target_net,
'optimizer': torch.optim.Adam([{'name':'net_params', 'params':target_net.parameters(), 'base_lr':0.00005}], betas=(0.9, 0.999), weight_decay=0.0005),
'lr_step_ratio': 0.5,
'lr_step_size': 5,
'supervise':{
'out': {'MSE_loss': {'obj': nn.MSELoss(size_average=True), 'factor':0.1, 'weight': 1.0}},
},
'metrics': {}
})
train_dataset = RandomCropPhoto2PhotoData({
'crop_size': hr_size,
'crop_stride': 2,
'data_root': DATASET_DIR,
'desc_file_path': os.path.join(DATASET_TXT_DIR, DATASET_ID, 'train.txt'),
})
valid_dataset = Photo2PhotoData({
'data_root': DATASET_DIR,
'desc_file_path': os.path.join(DATASET_TXT_DIR, DATASET_ID, 'val.txt'),
})
'psnr': {
'obj': PSNR()
}
}
},
'not_show_gradient':
True
})
# =================== dataset =====================================================================
train_dataset = RandomCropPhoto2PhotoData({
'crop_size':
50,
'crop_stride':
2,
'data_root':
DATASET_DIR,
'desc_file_path':
os.path.join(DATASET_TXT_DIR, 'train.txt'),
'is_rotated':
True
})
valid_dataset = RandomCropPhoto2PhotoData({
'crop_size':
50,
'crop_stride':
2,
'data_root':
DATASET_DIR,
'desc_file_path':
os.path.join(DATASET_TXT_DIR, 'val.txt'),