42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import numpy as np
|
|
import os
|
|
import ntpath
|
|
import time
|
|
import utils
|
|
from scipy.misc import imresize
|
|
from tensorboardX import SummaryWriter
|
|
|
|
|
|
class TFVisualizer():
|
|
def __init__(self, opt):
|
|
self.tf_visualizer = SummaryWriter(os.path.join(opt.logDir, opt.name))
|
|
self.opt = opt
|
|
self.saved = False
|
|
self.ncols = 4
|
|
|
|
self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
|
|
with open(self.log_name, "a") as log_file:
|
|
now = time.strftime("%c")
|
|
log_file.write('================ Training Loss (%s) ================\n' % now)
|
|
|
|
def reset(self):
|
|
self.saved = False
|
|
|
|
# |visuals|: dictionary of images to display or save
|
|
def display_current_results(self, visuals, iter_mark, epoch, save_result):
|
|
for label, image in visuals.items():
|
|
img_gid = utils.tensor2imgrid(image)
|
|
self.tf_visualizer.add_image(label, img_gid, iter_mark)
|
|
|
|
# losses: dictionary of error labels and values
|
|
def plot_current_losses(self, iter_mark, losses):
|
|
# for label, loss in losses.items():
|
|
# self.tf_visualizer.add_scalar(label, loss, iter_mark)
|
|
self.tf_visualizer.add_scalars('training loss', losses, iter_mark)
|
|
|
|
def print_logs(self, message):
|
|
print(message)
|
|
with open(self.log_name, "a") as log_file:
|
|
log_file.write('%s\n' % message)
|
|
|