changed UNETR and SwinUNETR configs; adjust SwinUNETR code

Author: zengyuy

configs/SNEMI/SNEMI-Affinity-SwinUNETR.yaml

@@ -14,7 +14,7 @@ MODEL:
   USE_CHECKPOINT: False
   SPATIAL_DIMS: 3
   DOWNSAMPLE: 'merging'
-  USE_V2: False
+  USE_V2: True
 DATASET:
   OUTPUT_PATH: outputs/SNEMI_SwinUNETR
 SOLVER:

configs/SNEMI/SNEMI-Affinity-UNETR.yaml

@@ -8,7 +8,7 @@ MODEL:
   UNETR_NUM_HEADS: 12
   POS_EMBED: 'conv'
   NORM_NAME: 'instance'
-  CONV_BLOCK: True
+  CONV_BLOCK: False
   RES_BLOCK: True
   UNETR_DROPOUT_RATE: 0.0
 DATASET:

connectomics/config/defaults.py

@@ -227,6 +227,9 @@
 _C.DATASET.REDUCE_LABEL = True
 
 # Padding size for the input volumes
+# Due to the center crop in the data augmentor, regions close to the volume
+# border will never be sampled. Therefore we pad the input volume. For
+# large-scale dataset there is no need for padding.
 _C.DATASET.PAD_SIZE = [2, 64, 64]
 _C.DATASET.PAD_MODE = 'reflect'  # reflect, constant, symmetric
 

connectomics/model/arch/swinunetr.py

@@ -146,8 +146,8 @@ def __init__(
 
         self.encoder3 = UnetrBasicBlock(
             spatial_dims=spatial_dims,
-            in_channels=2 * feature_size,
-            out_channels=2 * feature_size,
+            in_channels=feature_size * 2,
+            out_channels=feature_size * 2,
             kernel_size=3,
             stride=1,
             norm_name=norm_name,
@@ -156,8 +156,8 @@ def __init__(
 
         self.encoder4 = UnetrBasicBlock(
             spatial_dims=spatial_dims,
-            in_channels=4 * feature_size,
-            out_channels=4 * feature_size,
+            in_channels=feature_size * 4,
+            out_channels=feature_size * 4,
             kernel_size=3,
             stride=1,
             norm_name=norm_name,
@@ -166,8 +166,8 @@ def __init__(
 
         self.encoder5 = UnetrBasicBlock(
             spatial_dims=spatial_dims,
-            in_channels=8 * feature_size,
-            out_channels=8 * feature_size,
+            in_channels=feature_size * 8,
+            out_channels=feature_size * 8,
             kernel_size=3,
             stride=1,
             norm_name=norm_name,
@@ -176,8 +176,8 @@ def __init__(
 
         self.encoder6 = UnetrBasicBlock(
             spatial_dims=spatial_dims,
-            in_channels=16 * feature_size,
-            out_channels=16 * feature_size,
+            in_channels=feature_size * 16,
+            out_channels=feature_size * 16,
             kernel_size=3,
             stride=1,
             norm_name=norm_name,
@@ -186,8 +186,8 @@ def __init__(
 
         self.decoder5 = UnetrUpBlock(
             spatial_dims=spatial_dims,
-            in_channels=16 * feature_size,
-            out_channels=8 * feature_size,
+            in_channels=feature_size * 16,
+            out_channels=feature_size * 8,
             kernel_size=3,
             upsample_kernel_size=up_kernel_size[4],
             norm_name=norm_name,
@@ -495,7 +495,7 @@ def forward(self, x, mask):
         else:
             attn = self.softmax(attn)
 
-        attn = self.attn_drop(attn)
+        attn = self.attn_drop(attn).to(v.dtype)
         x = (attn @ v).transpose(1, 2).reshape(b, n, c)
         x = self.proj(x)
         x = self.proj_drop(x)
@@ -1005,7 +1005,7 @@ def __init__(
                 self.layers4.append(layer)
             if self.use_v2:
                 layerc = UnetrBasicBlock(
-                    spatial_dims=3,
+                    spatial_dims=spatial_dims,
                     in_channels=embed_dim * 2**i_layer,
                     out_channels=embed_dim * 2**i_layer,
                     kernel_size=3,