Merge pull request #760 from analysiscenter/improvements

Improvements

Author: AlexeyKozhevin

batchflow/models/torch/base.py

@@ -196,11 +196,13 @@ class TorchModel(BaseModel, ExtractionMixin, OptimalBatchSizeMixin, Visualizatio
     trainable : sequence, optional
         Names of model parts to train. Should be a subset of names in `order` and can be used to freeze parameters.
 
-    init_weights : callable, 'best_practice_resnet', or None
+    init_weights : callable, 'best_practice_resnet', tuple, sequence of them or None
         Model weights initialization.
         If None, then default initialization is used.
         If 'best_practice_resnet', then common used non-default initialization is used.
         If callable, then callable applied to each layer.
+        If tuple, then the first element should be of the types above, and the second defines a model part to apply on.
+        If sequence, then each element should be of the types defined above: applies all init functions sequentially.
 
         Examples:
 
@@ -212,6 +214,7 @@ def callable_init(module): # example of a callable for init
                     nn.kaiming_normal_(module.weight)
 
             config = {'init_weights': callable_init}
+        - ``{'init_weights': ('best_practice_resnet', 'body')}`` # applies only at `body` module
 
 
     # Shapes: optional
@@ -854,7 +857,8 @@ def build_model(self, inputs=None):
             inputs = self.make_placeholder_data(to_device=True)
 
         if 'model' not in self.config:
-            self.model = Network(inputs=inputs, config=self.config, device=self.device)
+            with torch.no_grad():
+                self.model = Network(inputs=inputs, config=self.config, device=self.device)
         else:
             self.model = self.config['model']
 
@@ -902,12 +906,20 @@ def initialize_weights(self):
             # Parse model weights initialization
             init_weights = init_weights if isinstance(init_weights, list) else [init_weights]
 
-            for init_weights_function in init_weights:
+            for init_weights_ in init_weights:
+                if isinstance(init_weights_, tuple) and len(init_weights_) == 2:
+                    init_weights_function, init_weights_module = init_weights_
+                else:
+                    init_weights_function, init_weights_module = init_weights_, None
+
                 if init_weights_function in {'resnet', 'classic'}:
                     init_weights_function = best_practice_resnet_init
 
                 # Actual weights initialization
-                self.model.apply(init_weights_function)
+                if init_weights_module is None:
+                    self.model.apply(init_weights_function)
+                else:
+                    getattr(self.model, init_weights_module).apply(init_weights_function)
 
 
     # Transfer to/from device(s)
@@ -1010,6 +1022,10 @@ def train(self, inputs, targets, outputs=None, mode='train', lock=True, profile=
                 with the same keys and requested tensors as values.
         lock : bool
             If True, then model, loss and gradient update operations are locked, thus allowing for multithreading.
+        mode : None, str or callable
+            If None, then does nothing.
+            If str, then identifies mode to put the model in: one of ``'train'`` or ``'eval'``.
+            If callable, then applied to the model directly.
         sync_frequency : int, bool or None
             If int, then how often to apply accumulated gradients to the weights.
             If True, then value from config is used.
@@ -1336,6 +1352,10 @@ def predict(self, inputs, targets=None, outputs=None, lock=True, microbatch_size
         amp : None or bool
             If None, then use amp setting from config.
             If bool, then overrides the amp setting for prediction.
+        mode : None, str or callable
+            If None, then does nothing.
+            If str, then identifies mode to put the model in: one of ``'train'`` or ``'eval'``.
+            If callable, then applied to the model directly.
         no_grad : bool
             Whether to disable gradient computation during model evaluation.
         transfer_from_device : bool
@@ -1476,11 +1496,15 @@ def __call__(self, inputs, targets=None, outputs='predictions', lock=True,
 
     # Common utilities for train and predict
     def set_model_mode(self, mode):
-        """ Set model mode to either train or eval. """
+        """ Set model mode to either train or eval. If provided with a callable, applies it to the model directly. """
         if mode in {'train', 'training'}:
             self.model.train()
         elif mode in {'eval', 'predict', 'inference'}:
             self.model.eval()
+        elif mode is None:
+            pass
+        elif callable(mode):
+            self.model.apply(mode)
         else:
             raise ValueError(f'Unknown model mode={mode}')
 

batchflow/named_expr.py

@@ -343,12 +343,15 @@ def get(self, **kwargs):
             a = eval_expr(self.a, _call=False, **kwargs)
         else:
             a = eval_expr(self.a, **kwargs)
-        b = eval_expr(self.b, **kwargs)
-        c = eval_expr(self.c, **kwargs)
+
         if self.op in UNARY_OPS:
             return OPERATIONS[self.op](a)
+
+        b = eval_expr(self.b, **kwargs)
         if self.op in BINARY_OPS:
             return OPERATIONS[self.op](a, b)
+
+        c = eval_expr(self.c, **kwargs)
         return OPERATIONS[self.op](a, b, c)
 
     def assign(self, value, **kwargs):
@@ -401,6 +404,58 @@ def __repr__(self):
         return 'Unknown expression'
 
 
+class IF(NamedExpression):
+    """ Select either ``true`` or ``false``, based on ``condition``.
+    Useful for simple variables that change along the run of a pipeline.
+
+    Examples
+    --------
+    Select model mode based on the current pipeline iteration::
+        mode = IF(condition=I.current<450, true='train', false='eval')
+
+    Train the last 20% with larger batch size::
+        batch_size = IF(condition=I.ratio > 0.8, true=256, false=128)
+
+    Notes
+    -----
+    An alternative to this named expression is to use ``F``::
+        def select_batch_size(ratio):
+            return 256 if ratio > 0.8 else 128
+
+        batch_size = F(select_batch_size)(ratio=I.ratio)
+
+    Or, with a lambda::
+        batch_size = F(lambda ratio: 256 if ratio > 0.8 else 128)(ratio=I.ratio)
+
+    ``F`` is recommended where more flexibility is needed, and ``IF`` can be used for simple binary choices.
+    """
+    def __init__(self, condition, true, false, mode='w', **kwargs):
+        super().__init__('#!__if__', mode=mode, **kwargs)
+        self.condition = condition
+        self.true = true
+        self.false = false
+
+    def get(self, **kwargs):
+        """ Select based on condition. """
+        condition = eval_expr(self.condition, **kwargs)
+
+        if bool(condition):
+            return eval_expr(self.true, **kwargs)
+        return eval_expr(self.false, **kwargs)
+
+
+    def assign(self, value, **kwargs):
+        """ Assign a value to a named expression, based on condition. """
+        _, kwargs = self._get_params(**kwargs)
+
+        condition = eval_expr(self.condition, **kwargs)
+
+        if bool(condition):
+            self.true.assign(value, **kwargs)
+        else:
+            self.false.assign(value, **kwargs)
+
+
 class B(NamedExpression):
     """ Batch component or attribute name
 

batchflow/notifier.py

@@ -93,6 +93,7 @@ class Notifier:
         If str, then either registered monitor identifiers or names of pipeline variables.
         Named expressions are evaluated with the pipeline.
         If callable, then it is used to retrieve the container with data.
+        Should accept variable named arguments in the signature and may return ``None`` value to disable its plotting.
         If sequence, then it is used as the container with data.
         If dict, then 'source' key should be one of the above to identify container.
         Other available keys:
@@ -181,7 +182,7 @@ def __init__(self, bar='a', disable=False, frequency=1, monitors=None, graphs=No
                 elif isinstance(source, str):
                     container['name'] = source
                 elif callable(source):
-                    container['name'] = '<unknown_callable>'
+                    container['name'] = source.__name__
                 else:
                     container['name'] = '<unknown_container>'
 
@@ -364,7 +365,7 @@ def update_data(self, pipeline=None, batch=None):
                 container['data'] = value
 
             elif callable(source):
-                container['data'] = source()
+                container['data'] = source(container=container, notifier=self, pipeline=pipeline, batch=batch)
 
             else:
                 raise TypeError(f'Unknown type of `source`, {type(source)}!')
@@ -382,7 +383,7 @@ def make_plotter(self, num_graphs=None, layout='horizontal', figsize=None, ncols
         """ Make canvas for plotting graphs. """
         from .plotter import plot
         if num_graphs is None:
-            num_graphs = len(self.data_containers)
+            num_graphs = sum(container['data'] is not None for container in self.data_containers)
 
         if ncols is None and nrows is None:
             if layout in ['h', 'horizontal']:
@@ -430,7 +431,8 @@ def update_plot(self, index=0, add_suptitle=False, savepath=None, clear_display=
             self.plotter.config['suptitle'] = self.bar.format_meter(**fmt)
             self.plotter.annotate()
 
-        for i, container in enumerate(self.data_containers):
+        data_containers = [container for container in self.data_containers if container['data'] is not None]
+        for i, container in enumerate(data_containers):
             if i >= index:
                 subplot_index = i - index
                 subplot_config = plot_config.maybe_index(subplot_index)
@@ -461,7 +463,7 @@ def update_subplot(self, container, index, **kwargs):
         plot_config = container.get('plot_config', {})
         plot_config = {**plot_config, **kwargs}
 
-        x = np.arange(len(data))
+        x = np.arange(len(data)) if hasattr(data, '__len__') else None
         y = data
         if self.slice not in [None, slice(None)]:
             x = np.array(x)[self.slice]
@@ -471,6 +473,8 @@ def update_subplot(self, container, index, **kwargs):
             plot_function(ax=subplot.ax, index=index, x=x, y=y, container=container, notifier=self, **plot_config)
         elif isinstance(source, ResourceMonitor):
             source.plot(plotter=self.plotter, positions=index, **plot_config)
+        elif data is None:
+            pass
         else:
             source_defaults = {'title': name}
             if isinstance(data, (tuple, list)) or (isinstance(data, np.ndarray) and data.ndim == 1):