|
import ntpath |
|
import os |
|
import sys |
|
import time |
|
from pathlib import Path |
|
from subprocess import PIPE, Popen |
|
|
|
import numpy as np |
|
|
|
from util import now_time |
|
from . import html, util |
|
|
|
|
|
try: |
|
import wandb |
|
except ImportError: |
|
print( |
|
'Warning: wandb package cannot be found. The option "--use_wandb" will result in error.' |
|
) |
|
if sys.version_info[0] == 2: |
|
VisdomExceptionBase = Exception |
|
else: |
|
VisdomExceptionBase = ConnectionError |
|
|
|
|
|
def save_images( |
|
webpage, visuals, image_path, aspect_ratio=1.0, width=256, use_wandb=False |
|
): |
|
"""Save images to the disk. |
|
|
|
Parameters: |
|
use_wandb: |
|
webpage (the HTML class) -- the HTML webpage class that stores these imaegs (see html.py for more details) |
|
visuals (OrderedDict) -- an ordered dictionary that stores (name, images (either tensor or numpy) ) pairs |
|
image_path (str) -- the string is used to create image paths |
|
aspect_ratio (float) -- the aspect ratio of saved images |
|
width (int) -- the images will be resized to width x width |
|
|
|
This function will save images stored in 'visuals' to the HTML file specified by 'webpage'. |
|
""" |
|
image_dir = webpage.get_image_dir() |
|
short_path = ntpath.basename(image_path[0]) |
|
name = os.path.splitext(short_path)[0] |
|
|
|
webpage.add_header(name) |
|
ims, txts, links = [], [], [] |
|
ims_dict = {} |
|
for label, im_data in visuals.items(): |
|
im = util.tensor2im(im_data) |
|
image_name = "%s_%s.png" % (name, label) |
|
save_path = Path(image_dir, image_name) |
|
util.save_image(im, save_path, aspect_ratio=aspect_ratio) |
|
ims.append(image_name) |
|
txts.append(label) |
|
links.append(image_name) |
|
if use_wandb: |
|
ims_dict[label] = wandb.Image(im) |
|
webpage.add_images(ims, txts, links, width=width) |
|
if use_wandb: |
|
wandb.log(ims_dict) |
|
|
|
|
|
class Visualizer: |
|
"""This class includes several functions that can display/save images and print/save logging information. |
|
|
|
It uses a Python library 'visdom' for display, and a Python library 'dominate' (wrapped in 'HTML') for creating HTML files with images. |
|
""" |
|
|
|
def __init__(self, opt): |
|
"""Initialize the Visualizer class |
|
|
|
Parameters: |
|
opt -- stores all the experiment flags; needs to be a subclass of BaseOptions |
|
Step 1: Cache the training/test options |
|
Step 2: connect to a visdom server |
|
Step 3: create an HTML object for saveing HTML filters |
|
Step 4: create a logging file to store training losses |
|
""" |
|
self.opt = opt |
|
self.display_id = opt.display_id |
|
self.use_html = opt.isTrain and not opt.no_html |
|
self.win_size = opt.display_winsize |
|
self.name = opt.name |
|
self.port = opt.display_port |
|
self.saved = False |
|
self.use_wandb = opt.use_wandb |
|
self.wandb_project_name = opt.wandb_project_name |
|
self.current_epoch = 0 |
|
self.ncols = opt.display_ncols |
|
|
|
if ( |
|
self.display_id > 0 |
|
): |
|
import visdom |
|
|
|
self.vis = visdom.Visdom( |
|
server=opt.display_server, port=opt.display_port, env=opt.display_env |
|
) |
|
if not self.vis.check_connection(): |
|
self.create_visdom_connections() |
|
if self.use_wandb: |
|
self.wandb_run = ( |
|
wandb.init(project=self.wandb_project_name, name=opt.name, config=opt) |
|
if not wandb.run |
|
else wandb.run |
|
) |
|
self.wandb_run._label(repo="CycleGAN") |
|
|
|
if self.use_html: |
|
self.web_dir = Path(opt.checkpoints_dir, opt.name, "web" + now_time()) |
|
self.img_dir = self.web_dir.joinpath("images") |
|
print("Create web directory %s..." % self.web_dir) |
|
util.mkdirs([self.web_dir, self.img_dir]) |
|
|
|
self.log_name = Path(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): |
|
"""Reset the self.saved status""" |
|
self.saved = False |
|
|
|
def create_visdom_connections(self): |
|
"""If the program could not connect to Visdom server, this function will start a new server at port < self.port >""" |
|
cmd = sys.executable + " -m visdom.server -p %d &>/dev/null &" % self.port |
|
print("\n\nCould not connect to Visdom server. \n Trying to start a server....") |
|
print("Command: %s" % cmd) |
|
Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE) |
|
|
|
def display_current_results(self, visuals, epoch, save_result): |
|
"""Display current results on visdom; save current results to an HTML file. |
|
|
|
Parameters: |
|
visuals (OrderedDict) - - dictionary of images to display or save |
|
epoch (int) - - the current epoch |
|
save_result (bool) - - if save the current results to an HTML file |
|
""" |
|
if self.display_id > 0: |
|
ncols = self.ncols |
|
if ncols > 0: |
|
ncols = min(ncols, len(visuals)) |
|
h, w = next(iter(visuals.values())).shape[:2] |
|
table_css = """<style> |
|
|
|
table {border-collapse: separate; border-spacing: 4px; white-space: nowrap; text-align: center} |
|
|
|
table td {width: % dpx; height: % dpx; padding: 4px; outline: 4px solid black} |
|
|
|
</style>""" % ( |
|
w, |
|
h, |
|
) |
|
|
|
title = self.name |
|
label_html = "" |
|
label_html_row = "" |
|
images = [] |
|
image_numpy = None |
|
idx = 0 |
|
for label, image in visuals.items(): |
|
image_numpy = util.tensor2im(image) |
|
label_html_row += "<td>%s</td>" % label |
|
images.append(image_numpy.transpose([2, 0, 1])) |
|
idx += 1 |
|
if idx % ncols == 0: |
|
label_html += "<tr>%s</tr>" % label_html_row |
|
label_html_row = "" |
|
white_image = np.ones_like(image_numpy.transpose([2, 0, 1])) * 255 |
|
while idx % ncols != 0: |
|
images.append(white_image) |
|
label_html_row += "<td></td>" |
|
idx += 1 |
|
if label_html_row != "": |
|
label_html += "<tr>%s</tr>" % label_html_row |
|
try: |
|
self.vis.images( |
|
images, |
|
nrow=ncols, |
|
win=self.display_id + 1, |
|
padding=2, |
|
opts=dict(title=title + " images"), |
|
) |
|
label_html = "<table>%s</table>" % label_html |
|
self.vis.text( |
|
table_css + label_html, |
|
win=self.display_id + 2, |
|
opts=dict(title=title + " labels"), |
|
) |
|
except VisdomExceptionBase: |
|
self.create_visdom_connections() |
|
else: |
|
idx = 1 |
|
try: |
|
for label, image in visuals.items(): |
|
image_numpy = util.tensor2im(image) |
|
self.vis.image( |
|
image_numpy.transpose([2, 0, 1]), |
|
opts=dict(title=label), |
|
win=self.display_id + idx, |
|
) |
|
idx += 1 |
|
except VisdomExceptionBase: |
|
self.create_visdom_connections() |
|
if self.use_wandb: |
|
columns = [key for key, _ in visuals.items()] |
|
columns.insert(0, "epoch") |
|
result_table = wandb.Table(columns=columns) |
|
table_row = [epoch] |
|
ims_dict = {} |
|
for label, image in visuals.items(): |
|
image_numpy = util.tensor2im(image) |
|
wandb_image = wandb.Image(image_numpy) |
|
table_row.append(wandb_image) |
|
ims_dict[label] = wandb_image |
|
self.wandb_run.log(ims_dict) |
|
if epoch != self.current_epoch: |
|
self.current_epoch = epoch |
|
result_table.add_data(*table_row) |
|
self.wandb_run.log({"Result": result_table}) |
|
if self.use_html and ( |
|
save_result or not self.saved |
|
): |
|
self.saved = True |
|
|
|
for label, image in visuals.items(): |
|
image_numpy = util.tensor2im(image) |
|
img_path = Path(self.img_dir, "epoch%.3d_%s.png" % (epoch, label)) |
|
util.save_image(image_numpy, img_path) |
|
|
|
webpage = html.HTML( |
|
self.web_dir, "Experiment name = %s" % self.name, refresh=1 |
|
) |
|
for n in range(epoch, 0, -1): |
|
webpage.add_header("epoch [%d]" % n) |
|
ims, txts, links = [], [], [] |
|
|
|
for label, image_numpy in visuals.items(): |
|
|
|
img_path = "epoch%.3d_%s.png" % (n, label) |
|
ims.append(img_path) |
|
txts.append(label) |
|
links.append(img_path) |
|
webpage.add_images(ims, txts, links, width=self.win_size) |
|
webpage.save() |
|
|
|
def plot_current_losses(self, epoch, counter_ratio, losses): |
|
"""display the current losses on visdom display: dictionary of error labels and values |
|
|
|
Parameters: |
|
epoch (int) -- current epoch |
|
counter_ratio (float) -- progress (percentage) in the current epoch, between 0 to 1 |
|
losses (OrderedDict) -- training losses stored in the format of (name, float) pairs |
|
""" |
|
if not hasattr(self, "plot_data"): |
|
self.plot_data = {"X": [], "Y": [], "legend": list(losses.keys())} |
|
self.plot_data["X"].append(epoch + counter_ratio) |
|
self.plot_data["Y"].append([losses[k] for k in self.plot_data["legend"]]) |
|
try: |
|
self.vis.line( |
|
X=np.stack( |
|
[np.array(self.plot_data["X"])] * len(self.plot_data["legend"]), 1 |
|
), |
|
Y=np.array(self.plot_data["Y"]), |
|
opts={ |
|
"title": f"{self.name} loss over time", |
|
"legend": self.plot_data["legend"], |
|
"xlabel": "epoch", |
|
"ylabel": "loss", |
|
}, |
|
win=self.display_id, |
|
) |
|
except VisdomExceptionBase: |
|
self.create_visdom_connections() |
|
if self.use_wandb: |
|
self.wandb_run.log(losses) |
|
|
|
|
|
def print_current_losses(self, epoch, iters, losses, t_comp, t_data): |
|
"""print current losses on console; also save the losses to the disk |
|
|
|
Parameters: |
|
epoch (int) -- current epoch |
|
iters (int) -- current training iteration during this epoch (reset to 0 at the end of every epoch) |
|
losses (OrderedDict) -- training losses stored in the format of (name, float) pairs |
|
t_comp (float) -- computational time per data point (normalized by batch_size) |
|
t_data (float) -- data loading time per data point (normalized by batch_size) |
|
""" |
|
message = f"(epoch: {epoch:>2d}, iters: {iters:>4d}, time: {t_comp:.3f}, data: {t_data:.3f})" |
|
for k, v in losses.items(): |
|
message += f" {k:s}: {v:.3f}" |
|
print(message) |
|
|
|
with open(self.log_name, "a") as log_file: |
|
log_file.write(f"{message:s}\n") |
|
|