diff --git a/cortex/_lib/__init__.py b/cortex/_lib/__init__.py index 8cbfa75..4bacc3a 100644 --- a/cortex/_lib/__init__.py +++ b/cortex/_lib/__init__.py @@ -9,7 +9,6 @@ from . import config, exp, log_utils, models from .parsing import DEFAULT_ARGS, parse_args, update_args -from .viz import init as viz_init from .utils import print_hypers @@ -75,6 +74,8 @@ def setup_experiment(args, model=None, testmode=False): """ + + exp.setup_visualization(args.visualization) def update_nested_dicts(from_d, to_d): for k, v in from_d.items(): if (k in to_d) and isinstance(to_d[k], dict): @@ -95,8 +96,15 @@ def update_nested_dicts(from_d, to_d): experiment_args = copy.deepcopy(DEFAULT_ARGS) update_args(experiment_args, exp.ARGS) - if not testmode and not args.noviz: - viz_init(config.CONFIG.viz) + + + if args.visualization == 'visdom': + from .viz import init as viz_init + if not testmode and args.visualization != 'off': + viz_init(config.CONFIG.viz) + + + def _expand_model_hypers(args, model): d = {} @@ -222,8 +230,15 @@ def reload(reload_path): exp.setup_out_dir(args.out_path, config.CONFIG.out_path, exp.NAME, clean=args.clean) - str = print_hypers(exp.ARGS, s='Final hyperparameters: ') + + if args.visualization == 'tensorboard': + from .tensorborad import init as tb_init + if not testmode: + tb_init(exp.OUT_DIRS['tb']) + + str = print_hypers(exp.ARGS, s='Final hyperparameters: ', mode=args.visualization) logger.info(str) model.push_hyperparameters(exp.ARGS['model']) + model.update_visualization(exp.VISUALIZATION) return model, reload_nets, args.lax_reload diff --git a/cortex/_lib/exp.py b/cortex/_lib/exp.py index cc82b56..4c754a0 100644 --- a/cortex/_lib/exp.py +++ b/cortex/_lib/exp.py @@ -28,6 +28,7 @@ INFO = {'name': NAME, 'epoch': 0} DEVICE = torch.device('cpu') DEVICE_IDS = None +VISUALIZATION = 'visdom' def _file_string(prefix: str = '') -> str: @@ -152,6 +153,7 @@ def setup_out_dir(out_path: str, global_out_path: str, name: str = None, clean: binary_dir = path.join(out_path, 'binaries') image_dir = path.join(out_path, 'images') + tb_dir = path.join(out_path, 'tb') if clean: logger.warning('Cleaning directory (cannot be undone)') @@ -159,17 +161,21 @@ def setup_out_dir(out_path: str, global_out_path: str, name: str = None, clean: rmtree(binary_dir) if path.isdir(image_dir): rmtree(image_dir) + if path.isdir(tb_dir): + rmtree(tb_dir) if not path.isdir(binary_dir): os.mkdir(binary_dir) if not path.isdir(image_dir): os.mkdir(image_dir) + if not path.isdir(tb_dir): + os.mkdir(tb_dir) logger.info('Setting out path to `{}`'.format(out_path)) logger.info('Logging to `{}`'.format(path.join(out_path, 'out.log'))) set_file_logger(path.join(out_path, 'out.log')) - OUT_DIRS.update(binary_dir=binary_dir, image_dir=image_dir) + OUT_DIRS.update(binary_dir=binary_dir, image_dir=image_dir, tb=tb_dir) def setup_device(device: [int] or int): @@ -183,3 +189,11 @@ def setup_device(device: [int] or int): DEVICE = torch.device(device) else: logger.info('Using CPU') + +def setup_visualization(visualization: str): + global VISUALIZATION + if visualization not in ['visdom', 'tensorboard', 'off']: + raise ValueError('Choose valid argument for visualisation') + + VISUALIZATION = visualization + logger.info('Visualization: {}'.format(visualization)) \ No newline at end of file diff --git a/cortex/_lib/models.py b/cortex/_lib/models.py index ab631fc..9b877ff 100644 --- a/cortex/_lib/models.py +++ b/cortex/_lib/models.py @@ -267,6 +267,7 @@ class ModelPluginBase(metaclass=PluginType): ''' # Global attributes that all models have access to. + _viz = VizHandler() _data = data.DATA_HANDLER _optimizers = optimizer.OPTIMIZERS @@ -337,6 +338,14 @@ def __setattr__(self, key, value): super().__setattr__(key, value) + def update_visualization(self, visualization): + + if visualization == 'visdom': + from .viz import VizHandler + elif visualization == 'tensorboard': + from .tensorborad import VizHandler + self._viz = VizHandler() + def add_models(self, **kwargs): for k, v in kwargs.items(): self.add_model(k, v) diff --git a/cortex/_lib/parsing.py b/cortex/_lib/parsing.py index 1d5701a..a4f22e2 100644 --- a/cortex/_lib/parsing.py +++ b/cortex/_lib/parsing.py @@ -206,7 +206,7 @@ def make_argument_parser() -> argparse.ArgumentParser: parser.add_argument('-v', '--verbosity', type=int, default=1, help='Verbosity of the logging. (0, 1, 2)') parser.add_argument('-d', '--device', type=int, nargs='+', default=0) - parser.add_argument('-V', '--noviz', default=False, action='store_true', help='No visualization.') + parser.add_argument('-V', '--visualization', default='visdom', type=str, help='options: visdom, tensorboard, off') return parser diff --git a/cortex/_lib/tensorborad.py b/cortex/_lib/tensorborad.py new file mode 100644 index 0000000..e59e409 --- /dev/null +++ b/cortex/_lib/tensorborad.py @@ -0,0 +1,455 @@ +""" +Visualization. +""" +import logging +from os import path + +import imageio +import numpy as np +from PIL import Image, ImageDraw +from tensorboardX import SummaryWriter + +from . import data, exp +from .utils import convert_to_numpy, compute_tsne +from .viz_utils import tile_raster_images +import matplotlib +import subprocess +from cortex._lib.config import _yes_no + +matplotlib.use('Agg') +from matplotlib import pylab as plt # noqa E402 +from torchvision.utils import make_grid + + + +logger = logging.getLogger('cortex.viz') +config_font = None +visualizer = None +_options = dict(label_names=None, is_caption=False, is_attribute=False) + +CHARS = ['_', '\n', ' ', '!', '"', '%', '&', "'", '(', ')', ',', '-', '.', '/', + '0', '1', '2', '3', '4', '5', '8', '9', ':', ';', '=', '?', '\\', '`', + 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', + 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '*', '*', + '*'] +CHAR_MAP = dict((i, CHARS[i]) for i in range(len(CHARS))) + + +def init(out_dir): + + # global visualizer, config_font, viz_process + # if viz_config is not None and ('server' in viz_config.keys() or + # 'port' in viz_config.keys()): + # server = viz_config.get('server', None) + # port = viz_config.get('port', 8097) + # logger.info('Using visdom version {}'.format(visdom.__version__)) + # visualizer = visdom.Visdom(server=server, port=port) + # if not visualizer.check_connection(): + # if _yes_no("No Visdom server runnning on the configured address. " + # "Do you want to start it?"): + # viz_bash_command = "python -m visdom.server" + # viz_process = subprocess.Popen(viz_bash_command.split()) + # logger.info('Using visdom server at {}({})'.format(server, port)) + # else: + # visualizer = None + # else: + # if _yes_no("Visdom configuration is not specified. Please run 'cortex setup' " + # "to configure Visdom server. Do you want to continue with " + # "the default address ? (localhost:8097)"): + # viz_bash_command = "python -m visdom.server" + # viz_process = subprocess.Popen(viz_bash_command.split()) + # visualizer = visdom.Visdom() + # logger.info('Using local visdom server') + # else: + # visualizer = None + # config_font = viz_config.get('font') + + + + global visualizer + visualizer = SummaryWriter(out_dir) + logger.info('Starting Tensorboard writer') + + + +def setup(label_names=None, is_caption=False, is_attribute=False, char_map=None): + """Sets up visualization arguments + + Args: + img: TODO + label_names: TODO + is_caption: TODO + is_attribute: TODO + char_map: TODO + + """ + global _options, CHAR_MAP + if label_names is not None: + _options['label_names'] = label_names + _options['is_caption'] = is_caption + _options['is_attribute'] = is_attribute + if is_caption and is_attribute: + raise ValueError('Cannot be both attribute and caption') + if char_map is not None: + CHAR_MAP = char_map + + +class VizHandler(): + def __init__(self): + self.clear() + self.output_dirs = exp.OUT_DIRS + self.prefix = exp._file_string('') + self.image_scale = (-1, 1) + + def clear(self): + self.images = {} + self.scatters = {} + self.histograms = {} + self.heatmaps = {} + + def add_image(self, im, name='image', labels=None): + visualizer.add_image(name, make_grid(im)) + + def add_histogram(self, hist, name='histogram'): + # if name in self.histograms: + # logger.warning('{} already added' + # ' to visualization.' + # ' Use the name kwarg' + # .format(name)) + hist = convert_to_numpy(hist) + # self.histograms[name] = hist + visualizer.add_histogram(name, hist) + + def add_heatmap(self, hm, name='heatmap'): + if name in self.heatmaps: + logger.warning('{} already' + ' added to visualization.' + ' Use the name kwarg' + .format(name)) + hm = convert_to_numpy(hm) + self.heatmaps[name] = hm + + def add_scatter(self, sc, labels=None, name='scatter'): + sc = convert_to_numpy(sc) + labels = convert_to_numpy(labels) + + self.scatters[name] = (sc, labels) + visualizer.add_embedding(sc, ) + + def show(self): + image_dir = self.output_dirs['image_dir'] + for i, (k, (im, labels)) in enumerate(self.images.items()): + if image_dir: + logger.debug('Saving images to {}'.format(image_dir)) + out_path = path.join( + image_dir, '{}_{}_image.png'.format(self.prefix, k)) + else: + out_path = None + + save_images(im, 8, 8, out_file=out_path, labels=labels, + max_samples=64, image_id=1 + i, caption=k) + + for i, (k, (sc, labels)) in enumerate(self.scatters.items()): + + if sc.shape[1] == 1: + raise ValueError('1D-scatter not supported') + elif sc.shape[1] > 2: + logger.info('Scatter greater than 2D. Performing TSNE to 2D') + sc = compute_tsne(sc) + + if image_dir: + logger.debug('Saving scatter to {}'.format(image_dir)) + out_path = path.join( + image_dir, '{}_{}_scatter.png'.format(self.prefix, k)) + else: + out_path = None + + save_scatter(sc, out_file=out_path, + labels=labels, image_id=i, + title=k) + + for i, (k, hist) in enumerate(self.histograms.items()): + if image_dir: + logger.debug('Saving histograms to {}'.format(image_dir)) + out_path = path.join( + image_dir, '{}_{}_histogram.png'.format(self.prefix, k)) + else: + out_path = None + save_hist(hist, out_file=out_path, hist_id=i) + + for i, (k, hm) in enumerate(self.heatmaps.items()): + if image_dir: + logger.debug('Saving heatmap to {}'.format(image_dir)) + out_path = path.join( + image_dir, '{}_{}_heatmap.png'.format(self.prefix, k)) + else: + out_path = None + save_heatmap(hm, out_file=out_path, image_id=i, title=k) + self.clear() + + +def plot(plot_updates, init=False, viz_test_only=False): + """Updates the plots for the reults. + + Takes the last value from the summary and appends this to the visdom plot. + + """ + def get_X_Y_legend(key, v_train, v_test): + Y = [v_train] + legend = [] + + if v_test is not None: + Y.append(v_test) + X = [range(len(v_train)), range(len(v_test))] + + legend.append('{} (train)'.format(key)) + legend.append('{} (test)'.format(key)) + else: + legend.append(key) + X = [range(len(v_train))] + + return X, Y, legend + + train_summary = exp.SUMMARY['train'] + test_summary = exp.SUMMARY['test'] + for k in train_summary.keys(): + + if viz_test_only and k != 'times': + if k in test_summary.keys(): + v_train = test_summary[k] + v_test = None + else: + continue + else: + v_train = train_summary[k] + v_test = test_summary[k] if k in test_summary.keys() else None + + if isinstance(v_train, dict): + Y = [] + X = [] + legend = [] + for k_ in v_train: + vt = v_test.get(k_) if v_test is not None else None + X_, Y_, legend_ = get_X_Y_legend(k_, v_train[k_], vt) + Y += Y_ + X += X_ + legend += legend_ + else: + X, Y, legend = get_X_Y_legend(k, v_train, v_test) + + if plot_updates: + label = 'Per {} updates'.format(plot_updates) + else: + label = 'Epochs' + opts = dict( + xlabel=label, + legend=legend, + ylabel=k, + title=k) + + X = np.array(X).transpose() + Y = np.array(Y).transpose() + + if Y.shape[-1] > 0: + visualizer.line( + Y=Y, + X=X, + env=exp.NAME, + opts=opts, + win='line_{}'.format(k) + ) + + +def save_text(labels, out_file=None, text_id=0, + caption='', step=None): + labels = np.argmax(labels, axis=-1) + char_map = _options['label_names'] + l_ = [''.join([char_map[j] for j in label]) for label in labels] + + logger.info('{}: {}'.format(caption, l_[0])) + # visualizer.text('\n'.join(l_), env=exp.NAME, + # win='text_{}'.format(text_id)) + + visualizer.add_text('text_{}'.format(text_id), '\n'.join(l_), global_step=step) + + if out_file is not None: + with open(out_file, 'w') as f: + for l__ in l_: + f.write(l__) + + +def save_images(images, num_x, num_y, out_file=None, labels=None, # noqa C901 + max_samples=None, margin_x=5, margin_y=5, image_id=0, + caption='', title='', step=None): + ''' + + Args: + images: + num_x: + num_y: + out_file: + labels: + max_samples: + margin_x: + margin_y: + image_id: + caption: + title: + + Returns: + + ''' + if labels is not None: + if isinstance(labels, (tuple, list)): + labels = zip(*labels) + if max_samples is not None: + images = images[:max_samples] + + if labels is not None: + if _options['is_caption']: + margin_x = 80 + margin_y = 80 + elif _options['is_attribute']: + margin_x = 25 + margin_y = 200 + elif _options['label_names'] is not None: + margin_x = 20 + margin_y = 25 + else: + margin_x = 5 + margin_y = 12 + + images = images * 255. + + dim_c, dim_x, dim_y = images.shape[-3:] + if dim_c == 1: + arr = tile_raster_images( + X=images, img_shape=(dim_x, dim_y), tile_shape=(num_x, num_y), + tile_spacing=(margin_y, margin_x), bottom_margin=margin_y) + fill = 255 + else: + arrs = [] + for c in range(dim_c): + arr = tile_raster_images( + X=images[:, c].copy(), img_shape=(dim_x, dim_y), + tile_shape=(num_x, num_y), + tile_spacing=(margin_y, margin_x), + bottom_margin=margin_y, right_margin=margin_x) + arrs.append(arr) + + arr = np.array(arrs).transpose(1, 2, 0) + fill = (255, 255, 255) + + im = Image.fromarray(arr) + + if labels is not None: + idr = ImageDraw.Draw(im) + for i, label in enumerate(labels): + x_ = (i % num_x) * (dim_x + margin_x) + y_ = (i // num_x) * (dim_y + margin_y) + dim_y + if _options['is_caption']: + l_ = ''.join([CHAR_MAP[j] for j in label + if CHAR_MAP[j] != '\n']) + # l__ = [CHAR_MAP[j] for j in label] + l_ = l_.strip() + if len(l_) == 0: + l_ = '' + if len(l_) > 30: + l_ = '\n'.join( + [l_[x:x + 30] for x in range(0, len(l_), 30)]) + elif _options['is_attribute']: + attribs = [j for j, a in enumerate(label) if a == 1] + l_ = '\n'.join(_options['label_names'][a] for a in attribs) + elif _options['label_names'] is not None: + l_ = _options['label_names'][label] + l_ = l_.replace('_', '\n') + else: + l_ = str(label) + idr.text((x_, y_), l_, fill=fill) + + arr = np.array(im) + if arr.ndim == 3: + arr = arr.transpose(2, 0, 1) + # visualizer.image(arr, opts=dict(title=title, caption=caption), + # win='image_{}'.format(image_id), + # env=exp.NAME) + + visualizer.add_images(tag='image_{}'.format(image_id), img_tensor = arr, global_step=step) + + if out_file: + im.save(out_file) + + +def save_heatmap(X, out_file=None, caption='', title='', image_id=0): + # visualizer.heatmap( + # X=X, + # opts=dict( + # title=title, + # caption=caption), + # win='heatmap_{}'.format(image_id), + # env=exp.NAME) + raise NotImplementedError + + +def save_scatter(points, out_file=None, labels=None, caption='', title='', + image_id=0): + # if labels is not None: + # Y = (labels + 1.5).astype(int) + # else: + # Y = None + + # names = data.DATA_HANDLER.get_label_names() + # Y = Y - min(Y) + 1 + # if len(names) != max(Y): + # names = ['{}'.format(i + 1) for i in range(max(Y))] + + # visualizer.scatter( + # X=points, + # Y=Y, + # opts=dict( + # title=title, + # caption=caption, + # legend=names, + # markersize=5), + # win='scatter_{}'.format(image_id), + # env=exp.NAME) + raise NotImplementedError + + +def save_movie(images, num_x, num_y, out_file=None, movie_id=0): + # if out_file is None: + # logger.warning('`out_file` not provided. Not saving.') + # else: + # images_ = [] + # for i, image in enumerate(images): + # dim_c, dim_x, dim_y = image.shape[-3:] + # image = image.reshape((num_x, num_y, dim_c, dim_x, dim_y)) + # image = image.transpose(0, 3, 1, 4, 2) + # image = image.reshape(num_x * dim_x, num_y * dim_y, dim_c) + # images_.append(image) + # imageio.mimsave(out_file, images_) + + # visualizer.video(videofile=out_file, env=exp.NAME, + # win='movie_{}'.format(movie_id)) + raise NotImplementedError + + +def save_hist(scores, out_file, hist_id=0): + s = list(scores.values()) + bins = np.linspace(np.min(np.array(s)), + np.max(np.array(s)), 100) + plt.clf() + for k, v in scores.items(): + plt.hist(v, bins, alpha=0.5, label=k) + plt.legend(loc='upper right') + if out_file: + plt.savefig(out_file) + hists = tuple(np.histogram(v, bins=bins)[0] for v in s) + X = np.column_stack(hists) + visualizer.stem( + X=X, Y=np.array([0.5 * (bins[i] + bins[i + 1]) for i in range(99)]), + opts=dict(legend=['Real', 'Fake']), win='hist_{}'.format(hist_id), + env=exp.NAME) + + visualizer.add_histogram('hist_{}'.format(hist_id), scores) + #raise NotImplementedError diff --git a/cortex/_lib/train.py b/cortex/_lib/train.py index 948ec1c..36740b7 100644 --- a/cortex/_lib/train.py +++ b/cortex/_lib/train.py @@ -378,15 +378,22 @@ def main_loop(model, epochs=500, archive_every=10, save_on_best=None, plot_updates: If set, plot is more fine-grained for updates. ''' - info = print_hypers(exp.ARGS, s='Model hyperparameters: ', visdom_mode=True) + info = print_hypers(exp.ARGS, s='Model hyperparameters: ', mode=exp.VISUALIZATION) logger.info('Starting main loop.') - if visdom_off: + if exp.VISUALIZATION == 'off': viz.visualizer = None - - if (viz.visualizer): + elif exp.VISUALIZATION == 'visdom': viz.visualizer.text(info, env=exp.NAME, win='info') + elif exp.VISUALIZATION == 'tensorboard': + from . import tensorborad as tb + tb.visualizer.add_text('info', info) + + + + + total_time = 0. if eval_only: train_results_ = test_epoch(model, None, data_mode=train_mode, @@ -453,10 +460,23 @@ def main_loop(model, epochs=500, archive_every=10, save_on_best=None, train_results_last_total = train_results_total test_results_last_total = test_results_total - if viz.visualizer: - plot(plot_updates, init=(epoch == first_epoch), viz_test_only=viz_test_only) - model.viz.show() - model.viz.clear() + + if exp.VISUALIZATION == 'visdom': + plot(plot_updates, init=(epoch == first_epoch), viz_test_only=viz_test_only) + model.viz.show() + model.viz.clear() + elif exp.VISUALIZATION == 'tensorboard': + losses = {} + for key in train_results_last_total.keys(): + if isinstance(train_results_last_total[key], dict): + for key2 in train_results_last_total[key].keys(): + losses['train_{}'.format(key2)] = train_results_last_total[key][key2] + losses['test_{}'.format(key2)] = test_results_last_total[key][key2] + else: + tb.visualizer.add_scalars('{}'.format(key), {'train': train_results_last_total[key], 'test': test_results_last_total[key]}, epoch) + + tb.visualizer.add_scalars('losses', losses, epoch) + if (archive_every and epoch % archive_every == 0): exp.save(model, prefix=epoch) diff --git a/cortex/_lib/utils.py b/cortex/_lib/utils.py index 1c7742f..265b0da 100644 --- a/cortex/_lib/utils.py +++ b/cortex/_lib/utils.py @@ -117,8 +117,11 @@ class bcolors: ENDC = '\033[0m' -def bold(s, visdom_mode=False): - if visdom_mode: +def bold(s, mode='visdom'): + if mode == 'tensorboard': + bold_char = '' + end_char = '' + elif mode == 'visdom': bold_char = '' end_char = '' else: @@ -127,8 +130,11 @@ def bold(s, visdom_mode=False): return bold_char + s + end_char -def underline(s, visdom_mode=False): - if visdom_mode: +def underline(s, mode='visdom'): + if mode == 'tensorboard': + ul_char = '' + end_char = '' + elif mode == 'visdom': ul_char = '' end_char = '' else: @@ -137,8 +143,8 @@ def underline(s, visdom_mode=False): return ul_char + s + end_char -def print_hypers(d, prefix=None, s='', visdom_mode=False, level=0): - if visdom_mode: +def print_hypers(d, prefix=None, s='', mode='visdom', level=0): + if mode == 'visdom': newline = '
' space = '  ' else: @@ -149,12 +155,12 @@ def print_hypers(d, prefix=None, s='', visdom_mode=False, level=0): s += '{}{}{}'.format(newline, space, prefix) if level == 0: spaces = space * 30 - s += underline('{}: {}'.format(k, spaces), visdom_mode=visdom_mode) + s += underline('{}: {}'.format(k, spaces), mode=mode) else: s += '{}: '.format(k) if isinstance(v, dict) and len(v) > 0: - s = print_hypers(v, prefix + space, s=s, visdom_mode=visdom_mode, level=level + 1) + s = print_hypers(v, prefix + space, s=s, mode=mode, level=level + 1) else: - s += bold('{}'.format(v), visdom_mode=visdom_mode) + s += bold('{}'.format(v), mode=mode) return s \ No newline at end of file diff --git a/cortex/_lib/viz.py b/cortex/_lib/viz.py index 1d1df77..fce1025 100644 --- a/cortex/_lib/viz.py +++ b/cortex/_lib/viz.py @@ -249,6 +249,7 @@ def get_X_Y_legend(key, v_train, v_test): X = np.array(X).transpose() Y = np.array(Y).transpose() + if Y.shape[-1] > 0: visualizer.line( Y=Y, diff --git a/setup.py b/setup.py index 95621fb..0e28f2c 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,8 @@ install_requirements = [ 'imageio', 'matplotlib==2.2.3', 'progressbar2', 'scipy', 'sklearn', 'visdom', - 'pyyaml', 'pathlib', 'sphinxcontrib-napoleon', 'nibabel', 'torch', 'torchvision' + 'pyyaml', 'pathlib', 'sphinxcontrib-napoleon', 'nibabel', 'torch', 'torchvision', + 'tensorflow', 'tensorboard', 'tensorboardX' ] setup(name='cortex',