parser.add_argument('--batch-size', '-b', type=int, default=1) parser.add_argument('--learning-rate', '-l', type=float, default=0.001) parser.add_argument('--epoch_num', '-e', type=int, default=80) parser.add_argument('--modeldir', '-md', type=str, default='res/') parser.add_argument('--batch-size-small', '-bs', type=int, default=1) if __name__ == "__main__": args = parser.parse_args() qf_fullname = 'matlab' + str(args.qf) training_set = utils.ImageDir(args.training_dir, args.training_cdir, args.training_dir1, args.training_cdir1, args.training_dir2, args.training_cdir2, preload=False, transform=tv.transforms.Compose([ utils.RandomCrop(56), utils.GenerateMultiscale(), utils.ToTorchTensor() ])) training_loader = torch.utils.data.DataLoader(training_set, batch_size=args.batch_size, num_workers=1, shuffle=True) val_set = utils.ImageDir(args.val_dir, args.val_cdir, preload=False, transform=tv.transforms.Compose([
type=str, default= '/mnt/hdd/compression-artifacts-becnmark/models/HRD4K/PRN/QP37/ckpt/300') parser.add_argument('--qf', '-q', type=int, default=10) # parser.add_argument('--output-dir', '-o', type=str) if __name__ == '__main__': args = parser.parse_args() qf_fullname = 'matlab' + str(args.qf) testing_set = utils.ImageDir( args.testing_dir, args.testing_cdir, preload=False, transform=tv.transforms.Compose([ # utils.ValCrop(448), utils.Align2(8), # utils.GenerateJPEGPair(args.training_dir,args.qf), utils.GenerateMultiscale(), utils.ToTorchTensor() ])) testing_loader = torch.utils.data.DataLoader(testing_set, batch_size=1, num_workers=1) # jnet = nn.DataParallel(model_prn.PRN(),device_ids=[1]).cuda() torch.cuda.set_device(1) jnet = model_prn.PRN().cuda() w = torch.load(args.checkpoint) new_w = collections.OrderedDict() for key, v in w.items():