def eval():
# Initialize Logger
logger.initLogger(args.debug)
# print iSeeBetter architecture
utils.printNetworkArch(netG=model, netD=None)
# load model
modelPath = os.path.join(args.model)
utils.loadPreTrainedModel(gpuMode=args.gpu_mode, model=model, modelPath=modelPath)
model.eval()
count = 0
avg_psnr_predicted = 0.0
for batch in testing_data_loader:
input, target, neigbor, flow, bicubic = batch[0], batch[1], batch[2], batch[3], batch[4]
with torch.no_grad():
if cuda:
input = Variable(input).cuda(gpus_list[0])
bicubic = Variable(bicubic).cuda(gpus_list[0])
neigbor = [Variable(j).cuda(gpus_list[0]) for j in neigbor]
flow = [Variable(j).cuda(gpus_list[0]).float() for j in flow]
else:
input = Variable(input).to(device=device, dtype=torch.float)
bicubic = Variable(bicubic).to(device=device, dtype=torch.float)
neigbor = [Variable(j).to(device=device, dtype=torch.float) for j in neigbor]
flow = [Variable(j).to(device=device, dtype=torch.float) for j in flow]
t0 = time.time()
if args.chop_forward:
with torch.no_grad():
prediction = chop_forward(input, neigbor, flow, model, args.upscale_factor)
else:
with torch.no_grad():
prediction = model(input, neigbor, flow)
if args.residual:
prediction = prediction + bicubic
t1 = time.time()
print("==> Processing: %s || Timer: %.4f sec." % (str(count), (t1 - t0)))
save_img(prediction.cpu().data, str(count), True)
save_img(target, str(count), False)
prediction = prediction.cpu()
prediction = prediction.data[0].numpy().astype(np.float32)
prediction = prediction*255.
target = target.squeeze().numpy().astype(np.float32)
target = target*255.
psnr_predicted = PSNR(prediction, target, shave_border=args.upscale_factor)
print("PSNR Predicted = ", psnr_predicted)
avg_psnr_predicted += psnr_predicted
count += 1
print("Avg PSNR Predicted = ", avg_psnr_predicted/count)
def eval():
# Initialize Logger
logger.initLogger(args.debug)
# print iSeeBetter architecture
utils.printNetworkArch(netG=model, netD=None)
# load model
modelPath = os.path.join(args.model)
utils.loadPreTrainedModel(gpuMode=args.gpu_mode,
model=model,
modelPath=modelPath)
model.eval()
count = 0
avg_psnr_predicted = 0.0
for batch in testing_data_loader:
#import pdb; pdb.set_trace()
input, target, neigbor, flow, bicubic = batch[0], batch[1], batch[
2], batch[3], batch[4]
with torch.no_grad():
if cuda:
input = Variable(input).cuda(gpus_list[0])
bicubic = Variable(bicubic).cuda(gpus_list[0])
neigbor = [Variable(j).cuda(gpus_list[0]) for j in neigbor]
flow = [Variable(j).cuda(gpus_list[0]).float() for j in flow]
else:
input = Variable(input).to(device=device, dtype=torch.float)
bicubic = Variable(bicubic).to(device=device,
dtype=torch.float)
neigbor = [
Variable(j).to(device=device, dtype=torch.float)
for j in neigbor
]
flow = [
Variable(j).to(device=device, dtype=torch.float)
for j in flow
]
t0 = time.time()
if args.chop_forward:
with torch.no_grad():
prediction = chop_forward(input, neigbor, flow, model,
args.upscale_factor)
else:
with torch.no_grad():
# unroll lists
neigbor0 = neigbor[0]
neigbor1 = neigbor[1]
neigbor2 = neigbor[2]
neigbor3 = neigbor[3]
neigbor4 = neigbor[4]
neigbor5 = neigbor[5]
flow0 = flow[0]
flow1 = flow[1]
flow2 = flow[2]
flow3 = flow[3]
flow4 = flow[4]
flow5 = flow[5]
#summary(model, (input, neigbor0, neigbor1, neigbor2, neigbor3, neigbor4, neigbor5, flow0, flow1, flow2, flow3, flow4, flow5))
print(model)
prediction = model(input, neigbor, flow)
# traced_model = torch.jit.trace(model, (input, neigbor0, neigbor1, neigbor2, neigbor3, neigbor4, neigbor5, flow0, flow1, flow2, flow3, flow4, flow5))
# mlmodel = ct.converters.convert(traced_model, inputs=\
# [ct.TensorType(shape=input.shape), ct.TensorType(shape=neigbor0.shape),
# ct.TensorType(shape=neigbor1.shape), ct.TensorType(shape=neigbor2.shape),
# ct.TensorType(shape=neigbor3.shape), ct.TensorType(shape=neigbor4.shape),
# ct.TensorType(shape=neigbor5.shape), ct.TensorType(shape=flow0.shape),
# ct.TensorType(shape=flow1.shape), ct.TensorType(shape=flow2.shape),
# ct.TensorType(shape=flow3.shape), ct.TensorType(shape=flow4.shape),
# ct.TensorType(shape=flow5.shape)])#[input, neigbor, flow])
# mlmodel.save('gen.mlmodel')
if args.residual:
prediction = prediction + bicubic
t1 = time.time()
print("==> Processing: %s || Timer: %.4f sec." % (str(count),
(t1 - t0)))
save_img(prediction.cpu().data, str(count), True)
save_img(target, str(count), False)
prediction = prediction.cpu()
prediction = prediction.data[0].numpy().astype(np.float32)
prediction = prediction * 255.
target = target.squeeze().numpy().astype(np.float32)
target = target * 255.
psnr_predicted = PSNR(prediction,
target,
shave_border=args.upscale_factor)
print("PSNR Predicted = ", psnr_predicted)
avg_psnr_predicted += psnr_predicted
count += 1
print("Avg PSNR Predicted = ", avg_psnr_predicted / count)
def main():
""" Lets begin the training process! """
args = parser.parse_args()
# Initialize Logger
logger.initLogger(args.debug)
# Load dataset
logger.info('==> Loading datasets')
# print(args.file_list)
# sys.exit()
train_set = get_training_set(args.data_dir, args.nFrames,
args.upscale_factor, args.data_augmentation,
args.file_list, args.other_dataset,
args.patch_size, args.future_frame)
training_data_loader = DataLoader(dataset=train_set,
num_workers=args.threads,
batch_size=args.batchSize,
shuffle=True)
# Use generator as RBPN
netG = RBPN(num_channels=3,
base_filter=256,
feat=64,
num_stages=3,
n_resblock=5,
nFrames=args.nFrames,
scale_factor=args.upscale_factor)
logger.info('# of Generator parameters: %s',
sum(param.numel() for param in netG.parameters()))
# Use DataParallel?
if args.useDataParallel:
gpus_list = range(args.gpus)
netG = torch.nn.DataParallel(netG, device_ids=gpus_list)
# Use discriminator from SRGAN
netD = Discriminator()
logger.info('# of Discriminator parameters: %s',
sum(param.numel() for param in netD.parameters()))
# Generator loss
generatorCriterion = nn.L1Loss() if not args.APITLoss else GeneratorLoss()
# Specify device
device = torch.device(
"cuda:0" if torch.cuda.is_available() and args.gpu_mode else "cpu")
if args.gpu_mode and torch.cuda.is_available():
utils.printCUDAStats()
netG.cuda()
netD.cuda()
netG.to(device)
netD.to(device)
generatorCriterion.cuda()
# Use Adam optimizer
optimizerG = optim.Adam(netG.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-8)
optimizerD = optim.Adam(netD.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-8)
if args.APITLoss:
logger.info(
"Generator Loss: Adversarial Loss + Perception Loss + Image Loss + TV Loss"
)
else:
logger.info("Generator Loss: L1 Loss")
# print iSeeBetter architecture
utils.printNetworkArch(netG, netD)
if args.pretrained:
modelPath = os.path.join(args.save_folder + args.pretrained_sr)
utils.loadPreTrainedModel(gpuMode=args.gpu_mode,
model=netG,
modelPath=modelPath)
# sys.exit()
for epoch in range(args.start_epoch, args.nEpochs + 1):
runningResults = trainModel(epoch, training_data_loader, netG, netD,
optimizerD, optimizerG, generatorCriterion,
device, args)
if (epoch + 1) % (args.snapshots) == 0:
saveModelParams(epoch, runningResults, netG, netD)