feat(vnet): convert dropout_prob to a tuple

monai/networks/nets/vnet.py

@@ -133,7 +133,7 @@ def __init__(
         out_channels: int,
         nconvs: int,
         act: tuple[str, dict] | str,
-        dropout_prob: float | None = None,
+        dropout_prob: tuple[float | None, float] = (None, 0.5),
         dropout_dim: int = 3,
     ):
         super().__init__()
@@ -144,8 +144,8 @@ def __init__(
 
         self.up_conv = conv_trans_type(in_channels, out_channels // 2, kernel_size=2, stride=2)
         self.bn1 = norm_type(out_channels // 2)
-        self.dropout = dropout_type(dropout_prob) if dropout_prob is not None else None
-        self.dropout2 = dropout_type(0.5)
+        self.dropout = dropout_type(dropout_prob[0]) if dropout_prob[0] is not None else None
+        self.dropout2 = dropout_type(dropout_prob[1])
         self.act_function1 = get_acti_layer(act, out_channels // 2)
         self.act_function2 = get_acti_layer(act, out_channels)
         self.ops = _make_nconv(spatial_dims, out_channels, nconvs, act)
@@ -224,7 +224,8 @@ def __init__(
         in_channels: int = 1,
         out_channels: int = 1,
         act: tuple[str, dict] | str = ("elu", {"inplace": True}),
-        dropout_prob: float = 0.5,
+        dropout_prob_down: float | None = 0.5,
+        dropout_prob_up: tuple[float | None, float] = (0.5, 0.5),
         dropout_dim: int = 3,
         bias: bool = False,
     ):
@@ -236,10 +237,10 @@ def __init__(
         self.in_tr = InputTransition(spatial_dims, in_channels, 16, act, bias=bias)
         self.down_tr32 = DownTransition(spatial_dims, 16, 1, act, bias=bias)
         self.down_tr64 = DownTransition(spatial_dims, 32, 2, act, bias=bias)
-        self.down_tr128 = DownTransition(spatial_dims, 64, 3, act, dropout_prob=dropout_prob, bias=bias)
-        self.down_tr256 = DownTransition(spatial_dims, 128, 2, act, dropout_prob=dropout_prob, bias=bias)
-        self.up_tr256 = UpTransition(spatial_dims, 256, 256, 2, act, dropout_prob=dropout_prob)
-        self.up_tr128 = UpTransition(spatial_dims, 256, 128, 2, act, dropout_prob=dropout_prob)
+        self.down_tr128 = DownTransition(spatial_dims, 64, 3, act, dropout_prob=dropout_prob_down, bias=bias)
+        self.down_tr256 = DownTransition(spatial_dims, 128, 2, act, dropout_prob=dropout_prob_down, bias=bias)
+        self.up_tr256 = UpTransition(spatial_dims, 256, 256, 2, act, dropout_prob=dropout_prob_up)
+        self.up_tr128 = UpTransition(spatial_dims, 256, 128, 2, act, dropout_prob=dropout_prob_up)
         self.up_tr64 = UpTransition(spatial_dims, 128, 64, 1, act)
         self.up_tr32 = UpTransition(spatial_dims, 64, 32, 1, act)
         self.out_tr = OutputTransition(spatial_dims, 32, out_channels, act, bias=bias)