dataset = data_loader.load_data()

dataset_size = len(data_loader)

print('#training images = %d' % dataset_size)

total_steps = 0

for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):

epoch_start_time = time.time()

epoch_iter = 0

for i, data in enumerate(dataset):

iter_start_time = time.time()

total_steps += opt.batchSize

epoch_iter += opt.batchSize

model.set_input(data)

model.optimize_parameters()

if total_steps % opt.display_freq == 0:

results = model.get_current_visuals()

psnrMetric = PSNR(results['Restored_Train'],results['Sharp_Train'])

print('PSNR on Train = %f' %

(psnrMetric))

visualizer.display_current_results(results,epoch)

if total_steps % opt.print_freq == 0:

errors = model.get_current_errors()

t = (time.time() - iter_start_time) / opt.batchSize

visualizer.print_current_errors(epoch, epoch_iter, errors, t)

if opt.display_id > 0:

visualizer.plot_current_errors(epoch, float(epoch_iter)/dataset_size, opt, errors)

if total_steps % opt.save_latest_freq == 0:

print('saving the latest model (epoch %d, total_steps %d)' %

(epoch, total_steps))

model.save('latest')

if epoch % opt.save_epoch_freq == 0:

print('saving the model at the end of epoch %d, iters %d' %

(epoch, total_steps))

model.save('latest')

model.save(epoch)

print('End of epoch %d / %d \t Time Taken: %d sec' %

(epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))

if epoch > opt.niter: