Why does the output layer in DynUNet use ADN?

[lidialuq]

The last output layer in the DynUNet implementation uses a convolution + ADN (activation-dropout-normalization). The deep supervision outputs use only a convolution, which is what I would expect from the last layer too. Does anyone know why it's implemented this way?
Here's the output form torchsummary from a DynUNet with 2 deep supervision outputs:

===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
├─DynUNetSkipLayer: 1-1                       [-1, 64, 128, 128, 128]   --
├─UnetBasicBlock: 1-2                         [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetBasicBlock: 2-1                    [-1, 64, 128, 128, 128]   (recursive)
├─UnetBasicBlock: 1                           []                        --
|    └─Convolution: 2-2                       [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetBasicBlock: 2                      []                        --
|    |    └─Convolution: 3-1                  [-1, 64, 128, 128, 128]   (recursive)
├─UnetBasicBlock: 1                           []                        --
|    └─Convolution: 2                         []                        --
|    |    └─Conv3d: 3-2                       [-1, 64, 128, 128, 128]   6,912
|    |    └─ADN: 3-3                          [-1, 64, 128, 128, 128]   --
|    └─InstanceNorm3d: 2-3                    [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetBasicBlock: 2                      []                        --
|    |    └─InstanceNorm3d: 3-4               [-1, 64, 128, 128, 128]   --
├─UnetBasicBlock: 1                           []                        --
|    └─LeakyReLU: 2-4                         [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetBasicBlock: 2                      []                        --
|    |    └─LeakyReLU: 3-5                    [-1, 64, 128, 128, 128]   --
├─UnetBasicBlock: 1                           []                        --
|    └─Convolution: 2-5                       [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetBasicBlock: 2                      []                        --
|    |    └─Convolution: 3-6                  [-1, 64, 128, 128, 128]   (recursive)
├─UnetBasicBlock: 1                           []                        --
|    └─Convolution: 2                         []                        --
|    |    └─Conv3d: 3-7                       [-1, 64, 128, 128, 128]   110,592
|    |    └─ADN: 3-8                          [-1, 64, 128, 128, 128]   --
|    └─InstanceNorm3d: 2-6                    [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetBasicBlock: 2                      []                        --
|    |    └─InstanceNorm3d: 3-9               [-1, 64, 128, 128, 128]   --
├─UnetBasicBlock: 1                           []                        --
|    └─LeakyReLU: 2-7                         [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetBasicBlock: 2                      []                        --
|    |    └─LeakyReLU: 3-10                   [-1, 64, 128, 128, 128]   --
|    └─DynUNetSkipLayer: 2-8                  [-1, 96, 64, 64, 64]      --
├─ModuleList: 1                               []                        --
|    └─UnetBasicBlock: 2-9                    [-1, 96, 64, 64, 64]      --
├─DynUNetSkipLayer: 1                         []                        --
|    └─DynUNetSkipLayer: 2                    []                        --
|    |    └─UnetBasicBlock: 3-11              [-1, 96, 64, 64, 64]      (recursive)
├─ModuleList: 1                               []                        --
|    └─UnetBasicBlock: 2                      []                        --
|    |    └─Convolution: 3-12                 [-1, 96, 64, 64, 64]      165,888
|    |    └─InstanceNorm3d: 3-13              [-1, 96, 64, 64, 64]      --
|    |    └─LeakyReLU: 3-14                   [-1, 96, 64, 64, 64]      --
|    |    └─Convolution: 3-15                 [-1, 96, 64, 64, 64]      248,832
|    |    └─InstanceNorm3d: 3-16              [-1, 96, 64, 64, 64]      --
|    |    └─LeakyReLU: 3-17                   [-1, 96, 64, 64, 64]      --
├─DynUNetSkipLayer: 1                         []                        --
|    └─DynUNetSkipLayer: 2                    []                        --
|    |    └─DynUNetSkipLayer: 3-18            [-1, 128, 32, 32, 32]     49,066,242
├─ModuleList: 1                               []                        --
|    └─UnetBasicBlock: 2-10                   [-1, 128, 32, 32, 32]     --
|    |    └─Convolution: 3-19                 [-1, 128, 32, 32, 32]     331,776
|    |    └─InstanceNorm3d: 3-20              [-1, 128, 32, 32, 32]     --
|    |    └─LeakyReLU: 3-21                   [-1, 128, 32, 32, 32]     --
|    |    └─Convolution: 3-22                 [-1, 128, 32, 32, 32]     442,368
|    |    └─InstanceNorm3d: 3-23              [-1, 128, 32, 32, 32]     --
|    |    └─LeakyReLU: 3-24                   [-1, 128, 32, 32, 32]     --
|    └─UnetBasicBlock: 2-11                   [-1, 192, 16, 16, 16]     --
|    |    └─Convolution: 3-25                 [-1, 192, 16, 16, 16]     663,552
|    |    └─InstanceNorm3d: 3-26              [-1, 192, 16, 16, 16]     --
|    |    └─LeakyReLU: 3-27                   [-1, 192, 16, 16, 16]     --
|    |    └─Convolution: 3-28                 [-1, 192, 16, 16, 16]     995,328
|    |    └─InstanceNorm3d: 3-29              [-1, 192, 16, 16, 16]     --
|    |    └─LeakyReLU: 3-30                   [-1, 192, 16, 16, 16]     --
|    └─UnetBasicBlock: 2-12                   [-1, 256, 8, 8, 8]        --
|    |    └─Convolution: 3-31                 [-1, 256, 8, 8, 8]        1,327,104
|    |    └─InstanceNorm3d: 3-32              [-1, 256, 8, 8, 8]        --
|    |    └─LeakyReLU: 3-33                   [-1, 256, 8, 8, 8]        --
|    |    └─Convolution: 3-34                 [-1, 256, 8, 8, 8]        1,769,472
|    |    └─InstanceNorm3d: 3-35              [-1, 256, 8, 8, 8]        --
|    |    └─LeakyReLU: 3-36                   [-1, 256, 8, 8, 8]        --
|    └─UnetBasicBlock: 2-13                   [-1, 384, 4, 4, 4]        --
|    |    └─Convolution: 3-37                 [-1, 384, 4, 4, 4]        2,654,208
|    |    └─InstanceNorm3d: 3-38              [-1, 384, 4, 4, 4]        --
|    |    └─LeakyReLU: 3-39                   [-1, 384, 4, 4, 4]        --
|    |    └─Convolution: 3-40                 [-1, 384, 4, 4, 4]        3,981,312
|    |    └─InstanceNorm3d: 3-41              [-1, 384, 4, 4, 4]        --
|    |    └─LeakyReLU: 3-42                   [-1, 384, 4, 4, 4]        --
├─UnetBasicBlock: 1-3                         [-1, 512, 2, 2, 2]        --
|    └─Convolution: 2-14                      [-1, 512, 2, 2, 2]        --
|    |    └─Conv3d: 3-43                      [-1, 512, 2, 2, 2]        5,308,416
|    |    └─ADN: 3-44                         [-1, 512, 2, 2, 2]        --
|    └─InstanceNorm3d: 2-15                   [-1, 512, 2, 2, 2]        --
|    └─LeakyReLU: 2-16                        [-1, 512, 2, 2, 2]        --
|    └─Convolution: 2-17                      [-1, 512, 2, 2, 2]        --
|    |    └─Conv3d: 3-45                      [-1, 512, 2, 2, 2]        7,077,888
|    |    └─ADN: 3-46                         [-1, 512, 2, 2, 2]        --
|    └─InstanceNorm3d: 2-18                   [-1, 512, 2, 2, 2]        --
|    └─LeakyReLU: 2-19                        [-1, 512, 2, 2, 2]        --
├─ModuleList: 1                               []                        --
|    └─UnetUpBlock: 2-20                      [-1, 384, 4, 4, 4]        --
|    |    └─Convolution: 3-47                 [-1, 384, 4, 4, 4]        1,572,864
|    |    └─UnetBasicBlock: 3-48              [-1, 384, 4, 4, 4]        11,943,936
|    └─UnetUpBlock: 2-21                      [-1, 256, 8, 8, 8]        --
|    |    └─Convolution: 3-49                 [-1, 256, 8, 8, 8]        786,432
|    |    └─UnetBasicBlock: 3-50              [-1, 256, 8, 8, 8]        5,308,416
|    └─UnetUpBlock: 2-22                      [-1, 192, 16, 16, 16]     --
|    |    └─Convolution: 3-51                 [-1, 192, 16, 16, 16]     393,216
|    |    └─UnetBasicBlock: 3-52              [-1, 192, 16, 16, 16]     2,985,984
|    └─UnetUpBlock: 2-23                      [-1, 128, 32, 32, 32]     --
|    |    └─Convolution: 3-53                 [-1, 128, 32, 32, 32]     196,608
|    |    └─UnetBasicBlock: 3-54              [-1, 128, 32, 32, 32]     1,327,104
├─ModuleList: 1                               []                        --
|    └─UnetOutBlock: 2-24                     [-1, 2, 32, 32, 32]       --
|    |    └─Convolution: 3-55                 [-1, 2, 32, 32, 32]       258
├─ModuleList: 1                               []                        --
|    └─UnetUpBlock: 2-25                      [-1, 96, 64, 64, 64]      --
├─DynUNetSkipLayer: 1                         []                        --
|    └─DynUNetSkipLayer: 2                    []                        --
|    |    └─UnetUpBlock: 3-56                 [-1, 96, 64, 64, 64]      (recursive)
├─ModuleList: 1                               []                        --
|    └─UnetUpBlock: 2                         []                        --
|    |    └─Convolution: 3-57                 [-1, 96, 64, 64, 64]      98,304
|    |    └─UnetBasicBlock: 3-58              [-1, 96, 64, 64, 64]      746,496
├─ModuleList: 1                               []                        --
|    └─UnetOutBlock: 2-26                     [-1, 2, 64, 64, 64]       --
├─DynUNetSkipLayer: 1                         []                        --
|    └─DynUNetSkipLayer: 2                    []                        --
|    |    └─UnetOutBlock: 3-59                [-1, 2, 64, 64, 64]       (recursive)
├─ModuleList: 1                               []                        --
|    └─UnetOutBlock: 2                        []                        --
|    |    └─Convolution: 3-60                 [-1, 2, 64, 64, 64]       194
├─ModuleList: 1                               []                        --
|    └─UnetUpBlock: 2-27                      [-1, 64, 128, 128, 128]   --
├─DynUNetSkipLayer: 1                         []                        --
|    └─UnetUpBlock: 2-28                      [-1, 64, 128, 128, 128]   (recursive)
├─ModuleList: 1                               []                        --
|    |    └─Convolution: 3-61                 [-1, 64, 128, 128, 128]   49,152
├─DynUNetSkipLayer: 1                         []                        --
|    |    └─Convolution: 3-62                 [-1, 64, 128, 128, 128]   (recursive)
├─ModuleList: 1                               []                        --
|    |    └─UnetBasicBlock: 3-63              [-1, 64, 128, 128, 128]   331,776
├─DynUNetSkipLayer: 1                         []                        --
|    |    └─UnetBasicBlock: 3-64              [-1, 64, 128, 128, 128]   (recursive)
├─UnetOutBlock: 1-4                           [-1, 2, 128, 128, 128]    --
|    └─Convolution: 2-29                      [-1, 2, 128, 128, 128]    --
|    |    └─Conv3d: 3-65                      [-1, 2, 128, 128, 128]    130
|    |    └─ADN: 3-66                         [-1, 2, 128, 128, 128]    --
===============================================================================================
Total params: 99,890,760
Trainable params: 99,890,760
Non-trainable params: 0
Total mult-adds (G): 876.97
===============================================================================================
Input size (MB): 32.00
Forward/backward pass size (MB): 3802.12
Params size (MB): 381.05
Estimated Total Size (MB): 4215.18
===============================================================================================

[KumoLiu]

Hi @lidialuq, I think we didn't use act and norm in the deep_supervision_heads. It will add an "adn" module without activation and normalization when you set dropout. You could print the model directly to see the difference.

MONAI/monai/networks/blocks/dynunet_block.py

Lines 260 to 261 in ab800d8

	act=None,
	norm=None,

MONAI/monai/networks/blocks/convolutions.py

Lines 158 to 171 in ab800d8

	if act is None and norm is None and dropout is None:
	return
	self.add_module(
	"adn",
	ADN(
	ordering=adn_ordering,
	in_channels=out_channels,
	act=act,
	norm=norm,
	norm_dim=self.spatial_dims,
	dropout=dropout,
	dropout_dim=dropout_dim,
	),
	)

(deep_supervision_heads): ModuleList(
    (0): UnetOutBlock(
      (conv): Convolution(
        (conv): Conv3d(64, 2, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        (adn): ADN(
          (D): Dropout(p=0.1, inplace=False)
        )
      )
    )
  )

Hope it can help you, thanks!