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
 
 
@@ -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

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,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(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))

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)

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