def get_generator_model():
generator = GeneratorResNet(img_shape=img_shape,
res_blocks=residual_blocks,
c_dim=c_dim)
generator.load_state_dict(
torch.load(PATH_G, map_location=torch.device('cpu')))
generator.eval()
return generator
SCALE_FACTOR = opt.scale_factor
MODEL_NAME = opt.model_name
hr_shape = (opt.hr_height, opt.hr_width)
results = {'Test': {'psnr': [], 'ssim': []}}
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
generator = GeneratorResNet()
generator = nn.DataParallel(generator, device_ids=[0, 1, 2])
generator.to(device)
# generator.load_state_dict(torch.load("saved_models/generator_%d_%d.pth" % (4,99)))
generator.load_state_dict(torch.load("saved_models/" + MODEL_NAME))
generator.eval()
test_dataloader = DataLoader(
TestImageDataset("../My_dataset/single_channel_100000/%s" %
opt.test_dataset_name,
hr_shape=hr_shape,
scale_factor=opt.scale_factor), # change
batch_size=1,
shuffle=False,
num_workers=opt.n_cpu,
)
test_bar = tqdm(test_dataloader, desc='[testing datasets]')
test_out_path = 'testing_results/SRF_' + str(SCALE_FACTOR) + '/'
if not os.path.exists(test_out_path):