print('---------- Generator architecture -------------')
print_network(G)
print('---------- Discriminator architecture -------------')
print_network(D)
print('----------------------------------------------')
model_denoiser = os.path.join(opt.save_folder + 'VAE_denoiser.pth')
denoiser.load_state_dict(
torch.load(model_denoiser, map_location=lambda storage, loc: storage))
print('Pre-trained Denoiser model is loaded.')
if opt.pretrained:
model_G = os.path.join(opt.save_folder + opt.pretrained_sr)
model_D = os.path.join(opt.save_folder + opt.pretrained_D)
if os.path.exists(model_G):
G.load_state_dict(
torch.load(model_G, map_location=lambda storage, loc: storage))
print('Pre-trained Generator model is loaded.')
if os.path.exists(model_D):
D.load_state_dict(
torch.load(model_D, map_location=lambda storage, loc: storage))
print('Pre-trained Discriminator model is loaded.')
if cuda:
denoiser = denoiser.cuda(gpus_list[0])
G = G.cuda(gpus_list[0])
D = D.cuda(gpus_list[0])
HR_feat_extractor = HR_feat_extractor.cuda(gpus_list[0])
feat_extractor = feat_extractor.cuda(gpus_list[0])
L1_loss = L1_loss.cuda(gpus_list[0])
BCE_loss = BCE_loss.cuda(gpus_list[0])
Lap_loss = Lap_loss.cuda(gpus_list[0])
if os.path.exists(opt.model_denoiser):
# denoiser.load_state_dict(torch.load(opt.model_denoiser, map_location=lambda storage, loc: storage))
pretrained_dict = torch.load(opt.model_denoiser,
map_location=lambda storage, loc: storage)
model_dict = denoiser.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
denoiser.load_state_dict(model_dict)
print('Pre-trained Denoiser model is loaded.')
if os.path.exists(opt.model_SR):
model.load_state_dict(
torch.load(opt.model_SR, map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
def eval():
denoiser.eval()
model.eval()
LR_image = [
join(opt.input, x) for x in listdir(opt.input) if is_image_file(x)
]
SR_image = [
join(opt.Result, x) for x in listdir(opt.input) if is_image_file(x)
]