81 layer factory

monai/networks/blocks/convolutions.py

@@ -12,24 +12,43 @@
 import numpy as np
 import torch.nn as nn
 
-from monai.networks.layers.factories import get_conv_type, get_dropout_type, get_normalize_type
+from monai.networks.layers.factories import Dropout, Norm, Act, Conv, split_args
 from monai.networks.layers.convutils import same_padding
 
 
 class Convolution(nn.Sequential):
+    """
+    Constructs a convolution with optional dropout, normalization, and activation layers.
+    """
 
-    def __init__(self, dimensions, in_channels, out_channels, strides=1, kernel_size=3, instance_norm=True, dropout=0,
-                 dilation=1, bias=True, conv_only=False, is_transposed=False):
+    def __init__(self, dimensions, in_channels, out_channels, strides=1, kernel_size=3, act=Act.PRELU, 
+                 norm=Norm.INSTANCE, dropout=None, dilation=1, bias=True, conv_only=False, is_transposed=False):
         super().__init__()
         self.dimensions = dimensions
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.is_transposed = is_transposed
 
         padding = same_padding(kernel_size, dilation)
-        normalize_type = get_normalize_type(dimensions, instance_norm)
-        conv_type = get_conv_type(dimensions, is_transposed)
-        drop_type = get_dropout_type(dimensions)
+        conv_type = Conv[Conv.CONVTRANS if is_transposed else Conv.CONV, dimensions]
+
+        # define the normalisation type and the arguments to the constructor
+        norm_name, norm_args = split_args(norm)
+        norm_type = Norm[norm_name, dimensions]
+
+        # define the activation type and the arguments to the constructor
+        act_name, act_args = split_args(act)
+        act_type = Act[act_name]
+
+        if dropout:
+            # if dropout was specified simply as a p value, use default name and make a keyword map with the value
+            if isinstance(dropout, (int, float)):
+                drop_name = Dropout.DROPOUT
+                drop_args = {"p": dropout}
+            else:
+                drop_name, drop_args = split_args(dropout)
+
+            drop_type = Dropout[drop_name, dimensions]
 
         if is_transposed:
             conv = conv_type(in_channels, out_channels, kernel_size, strides, padding, strides - 1, 1, bias, dilation)
@@ -39,17 +58,16 @@ def __init__(self, dimensions, in_channels, out_channels, strides=1, kernel_size
         self.add_module("conv", conv)
 
         if not conv_only:
-            self.add_module("norm", normalize_type(out_channels))
-            if dropout > 0:  # omitting Dropout2d appears faster than relying on it short-circuiting when dropout==0
-                self.add_module("dropout", drop_type(dropout))
+            self.add_module("norm", norm_type(out_channels, **norm_args))
+            if dropout:
+                self.add_module("dropout", drop_type(**drop_args))
 
-            self.add_module("prelu", nn.modules.PReLU())
+            self.add_module("act", act_type(**act_args))
 
 
 class ResidualUnit(nn.Module):
-
-    def __init__(self, dimensions, in_channels, out_channels, strides=1, kernel_size=3, subunits=2, instance_norm=True,
-                 dropout=0, dilation=1, bias=True, last_conv_only=False):
+    def __init__(self, dimensions, in_channels, out_channels, strides=1, kernel_size=3, subunits=2, 
+                 act=Act.PRELU, norm=Norm.INSTANCE, dropout=None, dilation=1, bias=True, last_conv_only=False):
         super().__init__()
         self.dimensions = dimensions
         self.in_channels = in_channels
@@ -64,10 +82,13 @@ def __init__(self, dimensions, in_channels, out_channels, strides=1, kernel_size
 
         for su in range(subunits):
             conv_only = last_conv_only and su == (subunits - 1)
-            unit = Convolution(dimensions, schannels, out_channels, sstrides, kernel_size, instance_norm, dropout,
-                               dilation, bias, conv_only)
+            unit = Convolution(dimensions, schannels, out_channels, sstrides, 
+                               kernel_size, act, norm, dropout, dilation, bias, conv_only)
+
             self.conv.add_module("unit%i" % su, unit)
-            schannels = out_channels  # after first loop set channels and strides to what they should be for subsequent units
+
+            # after first loop set channels and strides to what they should be for subsequent units
+            schannels = out_channels  
             sstrides = 1
 
         # apply convolution to input to change number of output channels and size to match that coming from self.conv
@@ -79,7 +100,7 @@ def __init__(self, dimensions, in_channels, out_channels, strides=1, kernel_size
                 rkernel_size = 1
                 rpadding = 0
 
-            conv_type = get_conv_type(dimensions, False)
+            conv_type = Conv[Conv.CONV, dimensions]
             self.residual = conv_type(in_channels, out_channels, rkernel_size, strides, rpadding, bias=bias)
 
     def forward(self, x):

monai/networks/layers/factories.py

@@ -8,8 +8,169 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""
+Defines factories for creating layers in generic, extensible, and dimensionally independent ways. A separate factory
+object is created for each type of layer, and factory functions keyed to names are added to these objects. Whenever
+a layer is requested the factory name and any necessary arguments are passed to the factory object. The return value
+is typically a type but can be any callable producing a layer object. 
 
-from torch import nn as nn
+The factory objects contain functions keyed to names converted to upper case, these names can be referred to as members
+of the factory so that they can function as constant identifiers. eg. instance normalisation is named `Norm.INSTANCE`.
+
+For example, to get a transpose convolution layer the name is needed and then a dimension argument is provided which is 
+passed to the factory function:
+
+    dimension = 3
+    name = Conv.CONVTRANS
+    conv = Conv[name, dimension]
+
+This allows the `dimension` value to be set in constructor for example so that the dimensionality of a network is 
+parameterizable. Not all factories require arguments after the name, the caller must be aware which are required.
+"""
+
+from typing import Callable
+
+import torch.nn as nn
+
+
+class LayerFactory:
+    """
+    Factory object for creating layers, this uses given factory functions to actually produce the types or constructing
+    callables. These functions are referred to by name and can be added at any time.
+    """
+
+    def __init__(self):
+        self.factories = {}
+
+    @property
+    def names(self):
+        """
+        Produces all factory names.
+        """
+
+        return tuple(self.factories)
+
+    def add_factory_callable(self, name, func):
+        """
+        Add the factory function to this object under the given name.
+        """
+
+        self.factories[name.upper()] = func
+
+    def factory_function(self, name):
+        """
+        Decorator for adding a factory function with the given name.
+        """
+
+        def _add(func):
+            self.add_factory_callable(name, func)
+            return func
+
+        return _add
+
+    def get_constructor(self, factory_name, *args):
+        """
+        Get the constructor for the given factory name and arguments.
+        """
+
+        if not isinstance(factory_name, str):
+            raise ValueError("Factories must be selected by name")
+
+        fact = self.factories[factory_name.upper()]
+        return fact(*args)
+
+    def __getitem__(self, args):
+        """
+        Get the given name or name/arguments pair. If `args` is a callable it is assumed to be the constructor
+        itself and is returned, otherwise it should be the factory name or a pair containing the name and arguments.
+        """
+
+        # `args` is actually a type or constructor
+        if callable(args):
+            return args
+
+        # `args` is a factory name or a name with arguments
+        if isinstance(args, str):
+            name_obj, args = args, ()
+        else:
+            name_obj, *args = args
+
+        return self.get_constructor(name_obj, *args)
+
+    def __getattr__(self, key):
+        """
+        If `key` is a factory name, return it, otherwise behave as inherited. This allows referring to factory names
+        as if they were constants, eg. `Fact.FOO` for a factory Fact with factory function foo.
+        """
+
+        if key in self.factories:
+            return key
+
+        return super().__getattr__(key)
+
+
+def split_args(args):
+    """
+    Split arguments in a way to be suitable for using with the factory types. If `args` is a name it's interpreted 
+    """
+
+    if isinstance(args, str):
+        return args, {}
+    else:
+        name_obj, args = args
+
+        if not isinstance(name_obj, (str, Callable)) or not isinstance(args, dict):
+            raise ValueError(
+                "Layer specifiers must be single strings or pairs of the form (name/object-types, argument dict)"
+            )
+
+        return name_obj, args
+
+
+# Define factories for these layer types
+
+Dropout = LayerFactory()
+Norm = LayerFactory()
+Act = LayerFactory()
+Conv = LayerFactory()
+
+
+@Dropout.factory_function("dropout")
+def dropout_factory(dim):
+    types = [nn.Dropout, nn.Dropout2d, nn.Dropout3d]
+    return types[dim - 1]
+
+
+@Norm.factory_function("instance")
+def instance_factory(dim):
+    types = [nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d]
+    return types[dim - 1]
+
+
+@Norm.factory_function("batch")
+def batch_factory(dim):
+    types = [nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d]
+    return types[dim - 1]
+
+
+Act.add_factory_callable("relu", lambda: nn.modules.ReLU)
+Act.add_factory_callable("leakyrelu", lambda: nn.modules.LeakyReLU)
+Act.add_factory_callable("prelu", lambda: nn.modules.PReLU)
+
+
+@Conv.factory_function("conv")
+def conv_factory(dim):
+    types = [nn.Conv1d, nn.Conv2d, nn.Conv3d]
+    return types[dim - 1]
+
+
+@Conv.factory_function("convtrans")
+def convtrans_factory(dim):
+    types = [nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d]
+    return types[dim - 1]
+
+
+# old factory functions
 
 
 def get_conv_type(dim, is_transpose):

monai/networks/nets/unet.py

@@ -12,6 +12,7 @@
 import torch.nn as nn
 
 from monai.networks.blocks.convolutions import Convolution, ResidualUnit
+from monai.networks.layers.factories import Norm, Act
 from monai.networks.layers.simplelayers import SkipConnection
 from monai.networks.utils import predict_segmentation
 from monai.utils import export
@@ -23,7 +24,7 @@
 class UNet(nn.Module):
 
     def __init__(self, dimensions, in_channels, num_classes, channels, strides, kernel_size=3, up_kernel_size=3,
-                 num_res_units=0, instance_norm=True, dropout=0):
+                 num_res_units=0, act=Act.PRELU, norm=Norm.INSTANCE, dropout=0):
         super().__init__()
         assert len(channels) == (len(strides) + 1)
         self.dimensions = dimensions
@@ -34,7 +35,8 @@ def __init__(self, dimensions, in_channels, num_classes, channels, strides, kern
         self.kernel_size = kernel_size
         self.up_kernel_size = up_kernel_size
         self.num_res_units = num_res_units
-        self.instance_norm = instance_norm
+        self.act = act
+        self.norm = norm
         self.dropout = dropout
 
         def _create_block(inc, outc, channels, strides, is_top):
@@ -63,35 +65,21 @@ def _create_block(inc, outc, channels, strides, is_top):
     def _get_down_layer(self, in_channels, out_channels, strides, is_top):
         if self.num_res_units > 0:
             return ResidualUnit(self.dimensions, in_channels, out_channels, strides, self.kernel_size, self.num_res_units,
-                                self.instance_norm, self.dropout)
+                                self.act, self.norm, self.dropout)
         else:
-            return Convolution(self.dimensions, in_channels, out_channels, strides, self.kernel_size, self.instance_norm,
+            return Convolution(self.dimensions, in_channels, out_channels, strides, self.kernel_size, self.act, self.norm,
                                self.dropout)
 
     def _get_bottom_layer(self, in_channels, out_channels):
         return self._get_down_layer(in_channels, out_channels, 1, False)
 
     def _get_up_layer(self, in_channels, out_channels, strides, is_top):
-        conv = Convolution(self.dimensions,
-                           in_channels,
-                           out_channels,
-                           strides,
-                           self.up_kernel_size,
-                           self.instance_norm,
-                           self.dropout,
-                           conv_only=is_top and self.num_res_units == 0,
-                           is_transposed=True)
+        conv = Convolution(self.dimensions, in_channels, out_channels, strides, self.up_kernel_size, self.act, self.norm,
+                           self.dropout, conv_only=is_top and self.num_res_units == 0, is_transposed=True)
 
         if self.num_res_units > 0:
-            ru = ResidualUnit(self.dimensions,
-                              out_channels,
-                              out_channels,
-                              1,
-                              self.kernel_size,
-                              1,
-                              self.instance_norm,
-                              self.dropout,
-                              last_conv_only=is_top)
+            ru = ResidualUnit(self.dimensions, out_channels, out_channels, 1, self.kernel_size, 1, self.act, self.norm,
+                              self.dropout, last_conv_only=is_top)
             return nn.Sequential(conv, ru)
         else:
             return conv