def train(net, trainloader, device, optimizer, criterion, memory_back=False):
"""One epoch training of a network.
:param net: The given network.
:param trainloader: Pytorch DataLoader (train set)
:param device: Either 'cuda' or 'cpu'
:param optimizer: The used optimizer.
:param criterion: The loss function.
:param memory_back: Whether or not for each batch the memory of the optimizer (if it has one) should be
temporarily added back to the net's parameters to compute the several metrics
with these new parameters. It doesn't change the final net's parameters.
:return: (train_loss, train_acc, train_loss_with_memory_back, train_acc_with_memory_back,
L1/L2 norm ratio of the gradients, L1/L2 norm ratio of g)
"""
net.train()
train_loss = 0
mback_train_loss = 0
correct = 0
mback_correct = 0
total = 0
norm_ratio_val = 0
corrected_norm_ratio_val = 0
# -------------------------------------
# Loop for training
# -------------------------------------
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
if memory_back:
with TemporarilyAddMemory(optimizer): # import convex_opt
outputs = net(inputs)
loss = criterion(outputs, targets)
mback_train_loss += loss.item()
_, predicted = outputs.max(1)
mback_correct += predicted.eq(targets).sum().item()
# train using optimizer
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
norm_ratio_val += optimizer.gradient_norms_ratio()
corrected_norm_ratio_val += optimizer.corrected_gradient_norms_ratio()
train_loss += loss.item() # loss value
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item() # acc value
progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (train_loss / (batch_idx + 1), 100. * correct / total, correct, total))
loss = train_loss/(batch_idx + 1)
acc = 100. * correct/total
return loss, acc, mback_train_loss/(batch_idx + 1), 100. * mback_correct/total, norm_ratio_val/(batch_idx + 1), \
corrected_norm_ratio_val/(batch_idx + 1)
def get_inception_activations(imgs_path, batch_size):
print("Calculating Activation")
num_imgs = len(imgs_path)
inception_network = PartialInceptionNetwork()
inception_network = to_cuda(inception_network)
inception_network.eval()
n_batches = int(np.ceil(num_imgs / batch_size))
inception_activations = np.zeros((num_imgs, 2048), dtype=np.float32)
index = 0
while len(imgs_path) != 0:
imgs_path, new_batch_size, tensors = preprocess_imgs(
imgs_path, batch_size)
activations = inception_network(tensors)
activations = activations.detach().cpu().numpy()
assert activations.shape == (
tensors.shape[0],
2048), "Expexted output shape to be: {}, but was: {}".format(
(tensors.shape[0], 2048), activations.shape)
start_index = index * batch_size
end_index = start_index + new_batch_size
index += 1
inception_activations[start_index:end_index, :] = activations
current = num_imgs - len(imgs_path)
progress_bar(current, num_imgs, barLength=20)
print("--Done--")
return inception_activations
def generate_img(model, config, saved_path, show, generation_size, alpha=0.4, batch_size=32):
if not os.path.exists(saved_path):
os.makedirs(saved_path)
cnt = 0
for i in range(math.ceil(generation_size/batch_size)):
if i == math.ceil(generation_size/batch_size) - 1:
batch_size = generation_size - i*batch_size
test_samples_z = torch.randn(batch_size, config['z_dim'], dtype=torch.float32).to(device)
with torch.no_grad():
generated_images = model.generate(test_samples_z, final_resolution_idx=model.res_idx, alpha=alpha)
generated_images = generated_images * 0.5 + 0.5
# generated_images = inverse_normalize(tensor=generated_images, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
for tensor in generated_images:
cvimg = single_tensor_to_img(tensor, 128, 128)
image_name = str(cnt) + ".jpg"
cnt += 1
file_name = os.path.join(saved_path, image_name)
if show == True:
cv2.imshow("restored_image", cvimg)
cv2.waitKey(0)
cvimg = cvimg * 255
cv2.imwrite(file_name, cvimg)
progress_bar(i, math.ceil(generation_size/batch_size))
cv2.destroyAllWindows()
print("Finished")
def test(epoch):
global best_acc
net.eval()
test_loss = 0
correct = 0
total = 0
idx_list=[]
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = net(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (test_loss/(batch_idx+1), 100.*correct/total, correct, total))
# Save checkpoint.
#joblib.dump(idx_list,'masked_contour_correct_idx.z')
acc = 100.*correct/total
if acc > best_acc:
print('Saving..')
state = {
'net': net.state_dict(),
'acc': acc,
'epoch': epoch,
}
if not os.path.isdir('checkpoints'):
os.mkdir('checkpoints')
torch.save(state, pth_path)
best_acc = acc
def test(net, testloader, device, optimizer, criterion, memory_back=False):
"""One test evaluation of a network.
:param net: The given network.
:param testloader: Pytorch DataLoader (train set)
:param device: Either 'cuda' or 'cpu'
:param optimizer: The used optimizer.
:param criterion: The loss function.
:param memory_back: Whether or not for each batch the memory of the optimizer (if it has one) should be
temporarily added back to the net's parameters to compute the several metrics
with these new parameters. It doesn't change the final net's parameters.
:return: (train_loss, train_acc, train_loss_with_memory_back, train_acc_with_memory_back)
"""
net.eval()
test_loss = 0
mback_test_loss = 0
correct = 0
mback_correct = 0
total = 0
# -------------------------------------
# Loop for testing
# -------------------------------------
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
if memory_back:
with TemporarilyAddMemory(optimizer):
outputs = net(inputs)
loss = criterion(outputs, targets)
mback_test_loss += loss.item()
_, predicted = outputs.max(1)
mback_correct += predicted.eq(targets).sum().item()
outputs = net(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (test_loss / (batch_idx + 1), 100. * correct / total, correct, total))
loss = test_loss / (batch_idx + 1)
acc = 100. * correct / total
return loss, acc, mback_test_loss / (batch_idx + 1), 100. * mback_correct / total
def pareto_test_main(opt, logger, model, test_loader, export_images=False):
test_set_name = test_loader.dataset.opt['name']
logger.info('\nTesting [{:s}]...'.format(test_set_name))
idx = 0
avg_psnr = 0
avg_niqe = 0
for batch_num, val_data in enumerate(test_loader):
img_name = osp.splitext(osp.basename(val_data['LQ_path'][0]))[0]
dataset_dir = osp.join(opt['path']['results_root'], test_set_name)
util.mkdir(dataset_dir)
model.feed_data(val_data)
model.test()
visuals = model.get_current_visuals()
sr_img = util.tensor2img(visuals['SR']) # uint8
# ground truth image
# gt_img = util.tensor2img(visuals['GT']) # uint8
# Save SR images for reference
if export_images:
suffix = opt['suffix']
if suffix:
save_img_path = osp.join(dataset_dir,
img_name + suffix + '.png')
else:
save_img_path = osp.join(dataset_dir, img_name + '.png')
util.save_img(sr_img, save_img_path)
# calculate PSNR
item_psnr = util.tensor_psnr(model.real_H, model.fake_H)
if math.isfinite(item_psnr):
avg_psnr += item_psnr
idx += 1
# calculate NIQE
item_niqe = util.tensor_niqe(model.fake_H)
# item_niqe = 0
if math.isfinite(item_niqe):
avg_niqe += item_niqe
progress_bar(batch_num, len(test_loader), msg=None)
avg_psnr = avg_psnr / idx
avg_niqe = avg_niqe / idx
logger.info("PSNR {} NIQE {}".format(avg_psnr, avg_niqe))
return avg_psnr, avg_niqe
def train(epoch):
print('\nEpoch: %d' % epoch)
encoder.train()
decoder.train()
classifier.train()
train_loss1 = 0
train_loss2 = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer1.zero_grad()
optimizer2.zero_grad()
optimizer3.zero_grad()
feature = encoder(inputs)
picture = decoder(feature)
picture_ = picture.data.detach().cpu()
input_pic = torch.Tensor(
normalize(picture_, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5))).to(device)
input_pic.requires_grad = True
y = classifier(input_pic)
loss1 = criterion1(picture, inputs)
loss2 = criterion2(y, targets)
loss1.backward(retain_graph=True)
if attenuation != "nc":
loss2.backward()
picture.backward(input_pic.grad / attenuation)
optimizer1.step()
optimizer2.step()
optimizer3.step()
train_loss1 += loss1.item()
train_loss2 += loss2.item()
_, predicted = y.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(
batch_idx, len(trainloader),
'Loss1: %.3f | Loss2: %.3f | Acc: %.3f%% (%d/%d)' %
(train_loss1 / (batch_idx + 1), train_loss2 /
(batch_idx + 1), 100. * correct / total, correct, total))
def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
train_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (train_loss/(batch_idx+1), 100.*correct/total, correct, total))
predecessors_list: dict = {}
vv = set()
config.logger.info(config.__str__())
count_up = 0
count_down = config.depth_size
while True:
"""
This while True loop makes sure that we get all of the transitions from the transition generator.
When no more transitions are left, the buildin call "next" will throw a StopIteration exception.
This is catched, and we append a few labels to make atl_generation easier. (on the last state, and on the first)
"""
try:
count_up += 1
count_down -= 1
progress_bar(count_up, config.depth_size)
transition = next(transitions)
result['labeling'].append(transition[0])
result['transitions'].append(transition[1])
result['moves'].append(transition[2])
config.logger.info(f'Countdown: {count_down}')
for key, value in transition[3].items():
if predecessors_list.get(key) is None:
predecessors_list[key] = value
else:
[predecessors_list[key].add(x) for x in value]
except StopIteration:
result['labeling'][-1].append(10)
result['labeling'][0].append(11)
break
config.logger.info('done')
def test(epoch):
global best_loss
global save_df
encoder.eval()
decoder.eval()
classifier.eval()
train_loss1 = 0
train_loss2 = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(testloader):
# print_real_pic(inputs)
inputs, targets = inputs.to(device), targets.to(device)
feature = encoder(inputs)
picture = decoder(feature)
# print_comp_pic(picture.cpu())
# print('Done.')
# sleep(1000)
picture_ = picture.data.detach().cpu()
input_pic = torch.Tensor(
normalize(picture_, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5))).to(device)
y = classifier(input_pic)
loss1 = criterion1(picture, inputs)
loss2 = criterion2(y, targets)
train_loss1 += loss1.item()
train_loss2 += loss2.item()
_, predicted = y.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(
batch_idx, len(testloader),
'Loss1: %.6f | Loss2: %.4f | Acc: %.3f%% (%d/%d)' %
(train_loss1 / (batch_idx + 1), train_loss2 /
(batch_idx + 1), 100. * correct / total, correct, total))
if batch_idx == 585:
save_df = save_df.append(
pd.DataFrame(
{
'epoch': epoch,
'MSE_loss': train_loss1 / (batch_idx + 1),
'classifier_loss': train_loss2 / (batch_idx + 1),
'accuracy': correct / total
},
index=[0]))
if not os.path.isdir('dataframe'):
os.mkdir('dataframe')
save_df.to_csv('./dataframe/result_' + codir + '.csv', index=False)
# Save checkpoint.
acc = 100. * correct / total
if train_loss1 < best_loss:
print('Saving..')
state = {
'encoder': encoder.state_dict(),
'decoder': decoder.state_dict(),
'classifier': classifier.state_dict(),
'loss': train_loss1,
'epoch': epoch,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, './checkpoint/ckpt_' + codir + '.t7')
best_loss = train_loss1
def psnr_main(opt,
train_loader,
val_loader,
train_sampler,
logger,
resume_state=None,
tb_logger=None,
rank=-1):
# create model
model = create_model(opt)
# resume training
if resume_state:
logger.info('Resuming training from epoch: {}, iter: {}.'.format(
resume_state['epoch'], resume_state['iter']))
start_epoch = resume_state['epoch']
current_step = resume_state['iter']
model.resume_training(resume_state) # handle optimizers and schedulers
else:
current_step = 0
start_epoch = 0
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
total_epochs = int(opt['train']['nepochs'])
best_psnr = 0
patience = 0
all_results = []
for epoch in range(start_epoch, total_epochs):
if opt['dist']:
train_sampler.set_epoch(epoch)
for batch_num, train_data in enumerate(train_loader):
current_step += 1
# training
model.feed_data(train_data)
model.optimize_parameters(current_step)
progress_bar(batch_num, len(train_loader), msg=None)
# log
if epoch % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step, model.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
# batched validation
if epoch % opt['train']['val_freq'] == 0 and rank <= 0:
avg_psnr = 0.0
idx = 0
for batch_num, val_data in enumerate(val_loader):
img_name = os.path.splitext(
os.path.basename(val_data['LQ_path'][0]))[0]
img_dir = os.path.join(opt['path']['val_images'], img_name)
util.mkdir(img_dir)
model.feed_data(val_data)
model.test()
# calculate PSNR
item_psnr = util.tensor_psnr(model.real_H, model.fake_H)
if math.isfinite(item_psnr):
avg_psnr += item_psnr
idx += 1
progress_bar(batch_num, len(val_loader), msg=None)
avg_psnr = avg_psnr / idx
all_results.append(avg_psnr)
if avg_psnr < best_psnr:
patience += 1
if patience == opt['train']['epoch_patience']:
model.update_learning_rate(opt['train']['lr_decay'])
print(
"no improvement, final patience, updating learning rate to {}"
.format(model.get_current_learning_rate()))
patience = 0
else:
print("no improvement, patience {} out of {}".format(
patience, opt['train']['epoch_patience']))
if model.get_current_learning_rate() < opt['train']['min_lr']:
break
else:
best_psnr = avg_psnr
if rank <= 0:
logger.info('Saving models and training states.')
model.save('latest')
# log
logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr))
logger_val = logging.getLogger('val') # validation logger
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} (best: {:.4e})'.
format(epoch, current_step, avg_psnr, best_psnr))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
if rank <= 0:
logger.info('End of training.')
json.dump(all_results,
open(
os.path.join(opt['path']['log'],
'validation_results.json'), 'w'),
indent=2)
lpips_t0, lpips_t1 = self.vgg16(img, resize_images=True, return_lpips=True).chunk(2)
dist = (lpips_t0 - lpips_t1).square().sum(1) / self.epsilon ** 2
return dist
if __name__ == '__main__':
batch_size = 16
test_size = 100000
total_iterations = math.ceil(test_size/batch_size)
vgg = torch.jit.load("./score_calculator/vgg16.pt").eval().to(device)
model = StyleALAE(model_config=config, device=device)
model.load_train_state('./archived/FFHQ/StyleALAE-z-256_w-256_prog-(4,256)-(8,256)-(16,128)-(32,128)-(64,64)-(64,32)/checkpoints/ckpt_gs-120000_res-5=64x64_alpha-0.40.pt')
model.G.eval()
model.F.eval()
for space in ['w', 'z']:
for path_len in ['full','end']:
sampler = PPLNet(G=model.G, F=model.F, config=config, batch_size=batch_size, res_idx=model.res_idx, alpha=0.4, epsilon=0.0001, space=space, sampling=path_len, vgg16=vgg)
sampler.eval().requires_grad_(False).to(device)
dists = []
print("Calculating on ",space," latent space , path length is ", path_len )
for itr in range(total_iterations):
dist = sampler.forward()
dists.append(dist)
progress_bar(itr, total_iterations)
print("Finished")
dists = torch.cat(dists)[:test_size].cpu().numpy()
lo = np.percentile(dists, 1, interpolation='lower')
hi = np.percentile(dists, 99, interpolation='higher')
ppl = np.extract(np.logical_and(dists >= lo, dists <= hi), dists).mean()
print("The PPL is ", ppl)
def train_main(opt,
train_loader,
val_loader,
train_sampler,
logger,
resume_state=None,
tb_logger=None,
rank=-1):
# create model
model = create_model(opt)
try:
total_nfe = model.netG.module.conv_trunk.nfe
nfe = True
try:
nfe_count = json.load(
open(os.path.join(opt['path']['log'], "nfe_count.json")))
print("resuming NFE count from {}".format(
os.path.join(opt['path']['log'], "nfe_count.json")))
except FileNotFoundError:
print("no previous NFE count file found, starting from scratch")
nfe_count = []
except AttributeError:
nfe = False
total_nfe = None
nfe_count = None
best_niqe = 1e10
best_psnr = 0
patience = 0
# resume training
if resume_state:
logger.info(
'Resuming training from epoch: {}, iter: {}., psnr: {}, niqe: {}, patience: {}'
.format(resume_state['epoch'], resume_state['iter'],
resume_state['psnr'], resume_state['niqe'],
resume_state['patience']))
start_epoch = resume_state['epoch'] + 1
current_step = resume_state['iter']
best_psnr = resume_state.get('psnr', 0)
best_niqe = resume_state.get('niqe', 1e10)
patience = resume_state.get('patience', 0)
model.resume_training(resume_state) # handle optimizers and schedulers
else:
current_step = 0
start_epoch = 0
try:
if opt['train']['G_pretraining'] >= 1:
pretraining_epochs = opt['train']['G_pretraining']
else:
pretraining_epochs = 0
except (KeyError, TypeError):
pretraining_epochs = 0
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(
start_epoch, current_step))
total_epochs = int(opt['train']['nepochs'])
lr_decay = opt['train']['lr_decay']
min_lr = opt['train']['min_lr']
pretraining = False
all_results = []
start_time = time()
for epoch in range(start_epoch, total_epochs):
if pretraining_epochs > 0:
if epoch == 0:
pretraining = True
logger.info('Starting pretraining.')
if epoch == pretraining_epochs:
pretraining = False
logger.info(
'Pretraining done, adding feature and discriminator loss.')
if opt['dist']:
train_sampler.set_epoch(epoch)
if nfe:
epoch_nfe = []
for batch_num, train_data in enumerate(train_loader):
# try:
current_step += 1
# training
model.feed_data(train_data)
model.optimize_parameters(current_step, pretraining=pretraining)
if nfe:
last_nfe = model.netG.module.conv_trunk.nfe - total_nfe
total_nfe = model.netG.module.conv_trunk.nfe
epoch_nfe.append(last_nfe)
progress_bar(batch_num, len(train_loader), msg=None)
# except RuntimeError:
# continue
if nfe:
nfe_count.append(epoch_nfe)
# log
if epoch % opt['logger']['print_freq'] == 0:
logs = model.get_current_log()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
epoch, current_step, model.get_current_learning_rate())
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if rank <= 0:
tb_logger.add_scalar(k, v, current_step)
if rank <= 0:
logger.info(message)
# batched validation
if nfe:
epoch_nfe = []
if epoch % opt['train'][
'val_freq'] == 0 and rank <= 0 and epoch >= pretraining_epochs - 1:
avg_psnr = 0.0
avg_niqe = 0.0
idx = 0
for batch_num, val_data in enumerate(val_loader):
# try:
img_name = os.path.splitext(
os.path.basename(val_data['LQ_path'][0]))[0]
img_dir = os.path.join(opt['path']['val_images'], img_name)
util.mkdir(img_dir)
model.feed_data(val_data)
model.test()
if nfe:
last_nfe = model.netG.module.conv_trunk.nfe - total_nfe
total_nfe = model.netG.module.conv_trunk.nfe
epoch_nfe.append(last_nfe)
visuals = model.get_current_visuals()
sr_img = util.tensor2img(visuals['SR']) # uint8
# ground truth image
# gt_img = util.tensor2img(visuals['GT']) # uint8
# Save SR images for reference
save_img_path = os.path.join(
img_dir, '{:s}_{:d}.png'.format(img_name, current_step))
util.save_img(sr_img, save_img_path)
# calculate PSNR
item_psnr = util.tensor_psnr(model.real_H, model.fake_H)
if math.isfinite(item_psnr):
avg_psnr += item_psnr
idx += 1
# calculate NIQE
if opt['niqe']:
item_niqe = util.tensor_niqe(model.fake_H)
# item_niqe = 0
if math.isfinite(item_niqe):
avg_niqe += item_niqe
progress_bar(batch_num, len(val_loader), msg=None)
if nfe:
nfe_count.append(epoch_nfe)
json.dump(nfe_count,
open(
os.path.join(opt['path']['log'],
'nfe_count.json'), 'w'),
indent=2)
avg_psnr = avg_psnr / idx
avg_niqe = avg_niqe / idx
all_results.append((time() - start_time, avg_psnr, avg_niqe))
# save models and training states
if rank <= 0 and (avg_psnr > best_psnr
or avg_niqe < best_niqe - 10e-6):
logger.info('Saving models and training states.')
model.save(epoch)
model.save_training_state(epoch, current_step, best_psnr,
best_niqe, patience)
else:
patience += 1
if patience == opt['train']['epoch_patience']:
model.update_learning_rate(lr_decay)
print(
"no improvement, final patience, updating learning rate to {}"
.format(model.get_current_learning_rate()))
patience = 0
else:
print("no improvement, patience {} out of {}".format(
patience, opt['train']['epoch_patience']))
if model.get_current_learning_rate() < min_lr:
break
if avg_niqe < best_niqe:
best_niqe = avg_niqe
if avg_psnr > best_psnr:
best_psnr = avg_psnr
# log
logger.info('# Validation # PSNR: {:.4e} # NIQE: {:.4e}'.format(
avg_psnr, avg_niqe))
logger_val = logging.getLogger('val') # validation logger
logger_val.info(
'<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} niqe: {:.4e} (best: {:.4e}/{:.4e})'
.format(epoch, current_step, avg_psnr, avg_niqe, best_psnr,
best_niqe))
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
tb_logger.add_scalar('psnr', avg_psnr, current_step)
print('\n')
if rank <= 0:
# save results
logger.info('Saving the final model.')
model.save('latest')
logger.info('End of training.')
json.dump(all_results,
open(
os.path.join(opt['path']['log'],
'validation_results.json'), 'w'),
indent=2)
if nfe:
nfe_count.append(epoch_nfe)
json.dump(nfe_count,
open(os.path.join(opt['path']['log'], 'nfe_count.json'),
'w'),
indent=2)
# clear validation logger
logger_val.handlers.clear()
# print out graph of val psnr with time
fig, ax = plt.subplots()
y = list(zip(*all_results))
runtime, dev_psnr, dev_niqe = y[0], y[1], y[2]
ax.plot(runtime, dev_psnr, color='blue', label='Validation PSNR')
ax.set(xlabel='Time (s)', ylabel='Dev. PSNR.')
ax.legend(loc='upper right')
ax.grid()
plt.savefig(os.path.join(opt['path']['log'], "psnr_evolution.png"))