if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
test_set = get_eval_set(os.path.join(opt.input_dir, opt.test_dataset),
opt.upscale_factor)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
print('===> Building model')
if opt.model_type == 'DBPNLL':
model = DBPNLL(num_channels=3,
base_filter=64,
feat=256,
num_stages=10,
scale_factor=opt.upscale_factor) ###D-DBPN
elif opt.model_type == 'DBPN-RES-MR64-3':
model = DBPNITER(num_channels=3,
base_filter=64,
feat=256,
num_stages=3,
scale_factor=opt.upscale_factor) ###D-DBPN
else:
model = DBPN(num_channels=3,
base_filter=64,
feat=256,
num_stages=7,
scale_factor=opt.upscale_factor) ###D-DBPN
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
train_set = get_training_set(opt.data_dir, opt.hr_train_dataset,
opt.upscale_factor, opt.patch_size,
opt.data_augmentation)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
print('===> Building model ', opt.model_type)
if opt.model_type == 'DBPNLL':
model = DBPNLL(num_channels=3,
base_filter=64,
feat=256,
num_stages=10,
scale_factor=opt.upscale_factor)
#elif opt.model_type == 'DBPN-RES-MR64-3':
# model = DBPNITER(num_channels=3, base_filter=64, feat = 256, num_stages=3, scale_factor=opt.upscale_factor)
else:
model = DBPN(num_channels=3,
base_filter=64,
feat=256,
num_stages=7,
scale_factor=opt.upscale_factor)
model = torch.nn.DataParallel(model, device_ids=gpus_list)
###Discriminator
D = Discriminator(num_channels=3,
print('===> Loading datasets')
# train data loading
if opt.train_dataset.count('imagenet') > 0 or opt.hr_train_dataset.count('imagenet') > 1:
train_set = get_imagenet(opt.data_dir, opt.upscale_factor, opt.patch_size, opt.data_augmentation)
else:
train_set = get_training_set(opt.data_dir, opt.train_dataset, opt.hr_train_dataset, opt.upscale_factor, opt.patch_size, opt.data_augmentation)
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
# test data loading
test_set = get_test_set(opt.data_dir, opt.test_dataset, opt.hr_test_dataset, opt.upscale_factor, opt.patch_size)
testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False)
print('===> Building model ', opt.model_type)
if opt.model_type == 'DBPNLL':
model = DBPNLL(num_channels=3, base_filter=64, feat = 256, num_stages=10, scale_factor=opt.upscale_factor) ###For NTIRE2018
else:
model = DBPN(num_channels=3, base_filter=64, feat = 256, num_stages=7, scale_factor=opt.upscale_factor) ###D-DBPN
model = torch.nn.DataParallel(model, device_ids=gpus_list)
criterion = nn.L1Loss()
print('---------- Networks architecture -------------')
print_network(model)
print('----------------------------------------------')
if opt.pretrained:
model_name = os.path.join(opt.save_folder + opt.pretrained_sr)
if os.path.exists(model_name):
#model= torch.load(model_name, map_location=lambda storage, loc: storage)
model.load_state_dict(torch.load(model_name, map_location=lambda storage, loc: storage))
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
test_set = get_eval_set(os.path.join(opt.input_dir, opt.test_dataset),
opt.upscale_factor)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
print('===> Building model')
if opt.model_type == 'DBPNLL':
model = DBPNLL(num_channels=3,
base_filter=64,
feat=256,
num_stages=10,
scale_factor=opt.upscale_factor) ###D-DBPN
elif opt.model_type == 'DBPN-RES-MR64-3':
model = DBPNITER(num_channels=3,
base_filter=64,
feat=256,
num_stages=3,
scale_factor=opt.upscale_factor) ###D-DBPN
else:
model = TISHKOVETS_NET(num_channels=3, base_filter=64, feat=256)
if cuda:
model = torch.nn.DataParallel(model, device_ids=gpus_list)
# model.load_state_dict(torch.load(opt.model, map_location=lambda storage, loc: storage))
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
test_set = get_eval_set(os.path.join(opt.input_dir, opt.test_dataset),
opt.upscale_factor)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
print('===> Building model')
if opt.model_type == 'DBPNLL':
model = DBPNLL(num_channels=3,
base_filter=64,
feat=256,
num_stages=10,
scale_factor=opt.upscale_factor) ###D-DBPN
elif opt.model_type == 'DBPN-RES-MR64-3':
model = DBPNITER(num_channels=3,
base_filter=64,
feat=256,
num_stages=3,
scale_factor=opt.upscale_factor) ###D-DBPN
else:
model = DBPN(num_channels=3,
base_filter=64,
feat=256,
num_stages=7,
scale_factor=opt.upscale_factor) ###D-DBPN