Tensorboard option

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
 
 
@@ -95,8 +94,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)
+    exp.setup_visualization(args.visdom)
+
+    if args.visdom == 'viz':
+        from .viz import init as viz_init
+        if not testmode and not args.noviz:
+            viz_init(config.CONFIG.viz)
+
+
+
 
     def _expand_model_hypers(args, model):
         d = {}
@@ -222,6 +228,12 @@ def reload(reload_path):
         exp.setup_out_dir(args.out_path, config.CONFIG.out_path, exp.NAME,
                           clean=args.clean)
 
+
+    if args.visdom == 'tb':
+        from .tensorborad import init as tb_init
+        if not testmode and not args.noviz:
+            tb_init(exp.OUT_DIRS['tb'])
+
     str = print_hypers(exp.ARGS, s='Final hyperparameters: ')
     logger.info(str)
     model.push_hyperparameters(exp.ARGS['model'])

cortex/_lib/exp.py

@@ -28,6 +28,7 @@
 INFO = {'name': NAME, 'epoch': 0}
 DEVICE = torch.device('cpu')
 DEVICE_IDS = None
+VIS = 'viz'
 
 
 def _file_string(prefix: str = '') -> str:
@@ -152,24 +153,29 @@ 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)')
         if path.isdir(binary_dir):
             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(vis: str):
+    global VIS
+    if vis not in ['vis', 'tb']:
+        raise ValueError('Choose valid argument for visualisation')
+
+    VIS = vis
+    logger.info('Visualization: {}'.format(vis))

cortex/_lib/parsing.py

@@ -207,6 +207,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
                         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('-vis', '--visdom', default='vis', type=str, help='options: vis, tb')
     return parser