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))
print('Pre-trained SR model is loaded.')
if cuda:
model = model.cuda(gpus_list[0])
criterion = criterion.cuda(gpus_list[0])
optimizer = optim.Adam(model.parameters(),
lr=opt.lr,
betas=(0.9, 0.999),
eps=1e-8)
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
train(epoch)
#test()
# learning rate is decayed by a factor of 10 every half of total epochs
if (epoch + 1) % (opt.nEpochs / 2) == 0:
for param_group in optimizer.param_groups:
param_group['lr'] /= 10.0
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)
base_filter=192,
feat=48,
num_stages=3,
n_resblock=5,
nFrames=opt.nFrames,
scale_factor=opt.upscale_factor)
if cuda:
model = torch.nn.DataParallel(model, device_ids=[gpus_list[gpu_no]])
model.load_state_dict(
torch.load(opt.model, map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
if cuda:
model = model.cuda(gpus_list[gpu_no])
def eval():
model.eval()
count = 1
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():
input = Variable(input).cuda(gpus_list[gpu_no])
bicubic = Variable(bicubic).cuda(gpus_list[gpu_no])
neigbor = [Variable(j).cuda(gpus_list[gpu_no]) for j in neigbor]
flow = [Variable(j).cuda(gpus_list[gpu_no]).float() for j in flow]
'''
print('===> Building model ', opt.model_type)
if opt.model_type == 'RBPN':
model = RBPN(num_channels=3, base_filter=256, feat = 64, num_stages=3, n_resblock=5, nFrames=opt.nFrames, scale_factor=opt.upscale_factor)
elif opt.model_type == "TouNet":
model = TouNet(num_channels=3, base_filter=192, feat = 48, num_stages=3, n_resblock=5, nFrames=opt.nFrames, scale_factor=opt.upscale_factor)
if cuda:
model = torch.nn.DataParallel(model, device_ids=[gpus_list[gpu_num]])
model.load_state_dict(torch.load(opt.model, map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
if cuda:
model = model.cuda(gpus_list[gpu_num])
def eval(data_loader, video_name):
model.eval()
count=0
avg_psnr_predicted = 0.0
for batch in data_loader:
input, neigbor, flow, bicubic = batch[0], batch[1], batch[2], batch[3]
with torch.no_grad():
input = Variable(input).cuda(gpus_list[gpu_num])
bicubic = Variable(bicubic).cuda(gpus_list[gpu_num])
neigbor = [Variable(j).cuda(gpus_list[gpu_num]) for j in neigbor]
flow = [Variable(j).cuda(gpus_list[gpu_num]).float() for j in flow]
t0 = time.time()
