def get_GAN_AB_model(folder_model, model_name, device):
n_residual_blocks = 9 # this should be the same values used in training the G_AB model
G_AB = GeneratorResNet(input_shape=(3,0), num_residual_blocks = n_residual_blocks)
G_AB.load_state_dict(torch.load(folder_model + model_name, map_location=device ), )
if cuda:
G_AB = G_AB.to(device)
return G_AB
type=int,
default=8,
help="number of cpu threads to use during batch generation")
opt = parser.parse_args()
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,
)