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():