Repository: analysiscenter/batchflow

Issue:
Update torch.Vnet with EncoderDecoder
Do this after #345


Discussion:

Resolving patch:
diff --git a/batchflow/models/torch/__init__.py b/batchflow/models/torch/__init__.py
--- a/batchflow/models/torch/__init__.py
+++ b/batchflow/models/torch/__init__.py
@@ -8,6 +8,7 @@
                     SEResNeXt18, SEResNeXt34, SEResNeXt50, SEResNeXt101, SEResNeXt152
 from .densenet import DenseNet, DenseNetS, DenseNet121, DenseNet169, DenseNet201, DenseNet264
 from .unet import UNet, ResUNet, DenseUNet
+from .vnet import VNet
 from .densenet import SegmentationDenseNet, DenseNetFC56, DenseNetFC67, DenseNetFC103
 from .utils import get_shape, get_num_channels, get_num_dims, calc_padding, unpack_fn_from_config, safe_eval
 from .blocks import DefaultBlock, XceptionBlock, VGGBlock, ResBlock, DenseBlock
diff --git a/batchflow/models/torch/unet.py b/batchflow/models/torch/unet.py
--- a/batchflow/models/torch/unet.py
+++ b/batchflow/models/torch/unet.py
@@ -11,6 +11,15 @@ class UNet(EncoderDecoder):
 
     Parameters
     ----------
+    auto_build : dict, optional
+        Parameters for auto-building `filters` in accordance with the idea described in the original paper.
+        Note that any of `filters`, if passed, will be replaced by auto-built ones.
+
+        num_stages : int
+            number of encoder/decoder stages — defines network depth and the number of its skip connections
+        filters : int, optional
+            number of filters in first encoder block — each of the following ones will be doubled until embedding
+
     body : dict
         encoder : dict
             num_stages : int
@@ -71,20 +80,16 @@ def default_config(cls):
     def build_config(self):
         config = super().build_config()
 
-        num_stages = config.get('body/encoder/num_stages')
-
-        if config.get('body/encoder/blocks/filters') is None:
-            config['body/encoder/blocks/filters'] = [64 * 2**i for i in range(num_stages)]
-
-        if config.get('body/embedding/filters') is None:
-            config['body/embedding/filters'] = 64 * 2**num_stages
-
-        if config.get('body/decoder/blocks/filters') is None:
-            enc_filters = config.get('body/encoder/blocks/filters')
-            config['body/decoder/blocks/filters'] = enc_filters[::-1]
+        if config.get('auto_build'):
+            num_stages = config.get('auto_build/num_stages', 4)
+            filters = config.get('auto_build/filters', 64)
+            encoder_filters = [filters * 2**i for i in range(num_stages)]
 
-        if config.get('body/decoder/upsample/filters') is None:
-            config['body/decoder/upsample/filters'] = config.get('body/decoder/blocks/filters')
+            config['body/encoder/num_stages'] = num_stages
+            config['body/encoder/blocks/filters'] = encoder_filters
+            config['body/embedding/filters'] = encoder_filters[-1] * 2
+            config['body/decoder/num_stages'] = num_stages
+            config['body/decoder/blocks/filters'] = encoder_filters[::-1]
 
         return config
 
diff --git a/batchflow/models/torch/vnet.py b/batchflow/models/torch/vnet.py
new file mode 100644
--- /dev/null
+++ b/batchflow/models/torch/vnet.py
@@ -0,0 +1,106 @@
+"""  Milletari F. et al "`V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation
+<https://arxiv.org/abs/1606.04797>`_"
+"""
+from .encoder_decoder import EncoderDecoder
+from .blocks import ResBlock
+
+
+
+class VNet(EncoderDecoder):
+    """ VNet-like model.
+
+    Parameters
+    ----------
+    auto_build : dict, optional
+        Parameters for auto-building `filters` and `layout` in accordance with the idea described in the original paper.
+        Note that any of these config params, if passed, will be replaced by auto-built ones.
+
+        num_stages : int
+            number of encoder/decoder stages — defines network depth and the number of its skip connections
+        filters : int, optional
+            number of filters in first encoder block — each of the following ones will be doubled until embedding
+
+    body : dict
+        encoder : dict
+            num_stages : int
+                number of downsampling blocks (default=4)
+            blocks : dict
+                Parameters for pre-processing blocks:
+
+                filters : None, int, list of ints or list of lists of ints
+                    The number of filters in the output tensor.
+                    If int, same number of filters applies to all layers on all stages
+                    If list of ints, specifies number of filters in each layer of different stages
+                    If list of list of ints, specifies number of filters in different layers on different stages
+                    If not given or None, filters parameters in encoder/blocks, decoder/blocks and decoder/upsample
+                    default to values which make number of filters double
+                    on each stage of encoding and halve on each stage of decoding,
+                    provided that `decoder/skip` is `True`. Specify `filters=None` explicitly
+                    if you want to use custom `num_steps` and infer `filters`
+
+        decoder : dict
+            num_stages : int
+                number of upsampling blocks. Defaults to the number of downsamplings.
+
+            factor : None, int or list of ints
+                If int, the total upsampling factor for all stages combined.
+                If list, upsampling factors for each stage
+                If not given or None, defaults to [2]*num_stages
+
+            blocks : dict
+                Parameters for post-processing blocks:
+
+                filters : None, int, list of ints or list of lists of ints
+                    same as encoder/blocks/filters
+
+            upsample : dict
+                Parameters for upsampling (see :func:`~.layers.upsample`).
+
+                filters : int, list of ints or list of lists of ints
+                    same as encoder/blocks/filters
+
+    Notes
+    -----
+    For more parameters see :class:`~.EncoderDecoder`.
+    """
+    @classmethod
+    def default_config(cls):
+        config = super().default_config()
+
+        config['body/encoder/num_stages'] = 4
+        config['body/encoder/order'] = ['block', 'skip', 'downsampling']
+        config['body/encoder/blocks'] += dict(base=ResBlock, layout=['cna', 'cna'*2, 'cna'*3, 'cna'*3],
+                                              filters=[16, 32, 64, 128], kernel_size=5)
+        config['body/encoder/downsample'] += dict(layout='cna', filters=[32, 64, 128, 256], kernel_size=2, strides=2)
+
+        config['body/embedding'] += dict(base=ResBlock, layout='cna'*3, filters=256, kernel_size=5)
+
+        config['body/decoder/order'] = ['upsampling', 'combine', 'block']
+        config['body/decoder/blocks'] += dict(base=ResBlock, layout=['cna'*3, 'cna'*3, 'cna'*2, 'cna'],
+                                              filters=[256, 128, 64, 32], kernel_size=5)
+        config['body/decoder/upsample'] += dict(layout='tna', filters=[128, 64, 32, 16], kernel_size=2, strides=2)
+
+        config['loss'] = 'ce'
+        return config
+
+    def build_config(self):
+        config = super().build_config()
+
+        if config.get('auto_build'):
+            num_stages = config.get('auto_build/num_stages', 4)
+            filters = config.get('auto_build/filters', 16)
+            encoder_filters = [filters * 2**i for i in range(num_stages)]
+            encoder_layout = ['cna', 'cna'*2] + ['cna'*3] * (num_stages - 2) if num_stages != 1 else 'cna'
+            downsample_filters = [filters * 2**(i + 1) for i in range(num_stages)]
+
+            config['body/encoder/num_stages'] = num_stages
+            config['body/encoder/blocks/filters'] = encoder_filters
+            config['body/encoder/blocks/layout'] = encoder_layout
+            config['body/encoder/downsample/filters'] = downsample_filters
+            config['body/embedding/filters'] = downsample_filters[-1]
+            config['body/decoder/num_stages'] = num_stages
+            config['body/decoder/blocks/filters'] = downsample_filters[::-1]
+            config['body/decoder/blocks/layout'] = encoder_layout[::-1]
+            config['body/decoder/upsample/filters'] = encoder_filters[::-1]
+
+        return config


Test patch:
diff --git a/batchflow/tests/notebooks/torch_test.ipynb b/batchflow/tests/notebooks/torch_test.ipynb
--- a/batchflow/tests/notebooks/torch_test.ipynb
+++ b/batchflow/tests/notebooks/torch_test.ipynb
@@ -522,7 +522,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.3"
+   "version": "3.6.8"
   },
   "varInspector": {
    "cols": {