add adn unet (#4185)

Signed-off-by: Wenqi Li <[email protected]>

Co-authored-by: Eric Kerfoot <[email protected]>

Author: wyli

monai/networks/blocks/acti_norm.py

@@ -18,8 +18,8 @@
 
 class ADN(nn.Sequential):
     """
-    Constructs a sequential module of optional activation, dropout, and normalization layers
-    (with an arbitrary order)::
+    Constructs a sequential module of optional activation (A), dropout (D), and normalization (N) layers
+    with an arbitrary order::
 
         -- (Norm) -- (Dropout) -- (Acti) --
 

monai/networks/nets/unet.py

@@ -68,6 +68,8 @@ class UNet(nn.Module):
         bias: whether to have a bias term in convolution blocks. Defaults to True.
             According to `Performance Tuning Guide <https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html>`_,
             if a conv layer is directly followed by a batch norm layer, bias should be False.
+        adn_ordering: a string representing the ordering of activation (A), normalization (N), and dropout (D).
+            Defaults to "NDA". See also: :py:class:`monai.networks.blocks.ADN`.
 
     Examples::
 
@@ -122,6 +124,7 @@ def __init__(
         norm: Union[Tuple, str] = Norm.INSTANCE,
         dropout: float = 0.0,
         bias: bool = True,
+        adn_ordering: str = "NDA",
         dimensions: Optional[int] = None,
     ) -> None:
 
@@ -155,6 +158,7 @@ def __init__(
         self.norm = norm
         self.dropout = dropout
         self.bias = bias
+        self.adn_ordering = adn_ordering
 
         def _create_block(
             inc: int, outc: int, channels: Sequence[int], strides: Sequence[int], is_top: bool
@@ -229,6 +233,7 @@ def _get_down_layer(self, in_channels: int, out_channels: int, strides: int, is_
                 norm=self.norm,
                 dropout=self.dropout,
                 bias=self.bias,
+                adn_ordering=self.adn_ordering,
             )
             return mod
         mod = Convolution(
@@ -241,6 +246,7 @@ def _get_down_layer(self, in_channels: int, out_channels: int, strides: int, is_
             norm=self.norm,
             dropout=self.dropout,
             bias=self.bias,
+            adn_ordering=self.adn_ordering,
         )
         return mod
 
@@ -279,6 +285,7 @@ def _get_up_layer(self, in_channels: int, out_channels: int, strides: int, is_to
             bias=self.bias,
             conv_only=is_top and self.num_res_units == 0,
             is_transposed=True,
+            adn_ordering=self.adn_ordering,
         )
 
         if self.num_res_units > 0:
@@ -294,6 +301,7 @@ def _get_up_layer(self, in_channels: int, out_channels: int, strides: int, is_to
                 dropout=self.dropout,
                 bias=self.bias,
                 last_conv_only=is_top,
+                adn_ordering=self.adn_ordering,
             )
             conv = nn.Sequential(conv, ru)
 

tests/test_unet.py

@@ -96,6 +96,7 @@
         "strides": (2, 2),
         "num_res_units": 1,
         "act": (Act.LEAKYRELU, {"negative_slope": 0.2}),
+        "adn_ordering": "NA",
     },
     (16, 4, 32, 64, 48),
     (16, 3, 32, 64, 48),