Swin UNETR | Flow Arena

Variants

SwinUNETR3D

64.1M

3D version with volumetric shifted window attention. Requires all spatial dims divisible by 32.

Attention: Shifted window self-attention (3D)

Window: 7

SwinUNETR2D

6.8M

2D version for Cartesian grid datasets. Requires spatial dims divisible by 32.

Attention: Shifted window self-attention (2D)

Window: 7

Exact params: 6,787,538

Benchmarks:

Two-Phase Oil-Water CO2 Radial

Training Provenance

These values are read from the generated benchmark result metadata, not from static method-page copy. When a model is retrained and the website data is regenerated, this section updates with the result records.

Arena training recipe

Epochs: 200

Batch size: 32

Learning rate: 0.0001

Loss: MSEGradientLoss, CombinedGradientLoss

Arena feature sets

Flow Arena input channels vary by fault complexity. Current result records for this method include 3 feature-set definitions.

PORO PERMX SWAT_0 PRESSURE_0 MULTX MULTY DEPTH Z_MLGRID

Strengths & Weaknesses

Strengths

Self-attention captures long-range spatial dependencies that CNNs may miss
Hierarchical multi-scale feature extraction via shifted window partitioning
Strong performance on datasets with large spatial extent and complex pressure fields
Pre-training potential — transformer weights can be initialized from large-scale vision pre-training
MONAI implementation is well-tested and production-grade

Weaknesses

Largest parameter count (64.1M 3D) — slowest to train and most memory-intensive
Strict divisibility-by-32 requirement limits applicability to arbitrary grid sizes
Cannot be applied to CO2 Nested dataset without architectural modifications
Transformer overhead may not be justified on smaller datasets where CNNs suffice
Window-based attention may still miss very long-range interactions that span multiple windows
Requires careful learning rate warmup and longer training schedules for convergence

Constraints

All spatial dimensions must be divisible by 32 (MONAI implementation requirement)
Excluded from CO2 Nested benchmark (depth dimension = 5 is not divisible by 32)
Largest model in the benchmark — requires more GPU memory and longer training time

Results Across Benchmarks

All results for Swin UNETR variants across all benchmark families. Sorted by the first field's rel-L² error by default.

Flow Arena

Property

Fault

Control

Loss

24 results

					Sat (frac)			Pres (bar)
#	Dataset	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	geostat_nf_rate	SwinUNETR	AbsLp(p=2)	200	0.1881	0.3405	0.046272	0.0370	0.0310	6.744236	16.3M
2	geostat_nf_bhp	SwinUNETR	AbsLp(p=2)	200	0.1882	0.3533	0.044299	0.0322	0.0280	6.993268	16.3M
3	geostat_nf_bhp	SwinUNETR	MSE	200	0.2008	0.3924	0.046215	0.0332	0.0285	7.094884	16.3M
4	geostat_nf_rate	SwinUNETR	MSE	200	0.2012	0.3854	0.042536	0.0414	0.0334	7.141094	16.3M
5	geostat_zt_rate	SwinUNETR	AbsLp(p=2)	200	0.2521	0.4546	0.059283	0.0684	0.0534	12.779385	16.3M
6	channels_nf_bhp	SwinUNETR	AbsLp(p=2)	200	0.2589	0.5225	0.069507	0.0416	0.0340	8.486210	16.3M
7	geostat_zt_rate	SwinUNETR	MSE	200	0.2601	0.4722	0.055375	0.0721	0.0570	13.083779	16.3M
8	geostat_vt_rate	SwinUNETR	AbsLp(p=2)	200	0.2664	0.5028	0.071910	0.0672	0.0519	11.766338	36.2M
9	geostat_zt_bhp	SwinUNETR	AbsLp(p=2)	200	0.2710	0.4854	0.060074	0.0481	0.0401	9.924505	16.3M
10	channels_nf_bhp	SwinUNETR	MSE	200	0.2712	0.5885	0.070083	0.0452	0.0370	9.148697	16.3M
11	geostat_vt_bhp	SwinUNETR	AbsLp(p=2)	200	0.2749	0.5230	0.072443	0.0458	0.0375	9.244534	36.2M
12	geostat_zt_bhp	SwinUNETR	MSE	200	0.2750	0.4835	0.058076	0.0485	0.0404	10.022796	16.3M
13	geostat_vt_rate	SwinUNETR	MSE	200	0.2802	0.5160	0.074990	0.0684	0.0527	12.038026	36.2M
14	geostat_vt_bhp	SwinUNETR	MSE	200	0.2809	0.5043	0.066116	0.0473	0.0383	9.439110	36.2M
15	channels_nf_rate	SwinUNETR	AbsLp(p=2)	200	0.2858	0.5270	0.069151	0.0529	0.0413	9.128984	16.3M
16	channels_nf_rate	SwinUNETR	MSE	200	0.2944	0.5695	0.066029	0.0553	0.0429	9.342718	16.3M
17	channels_vt_bhp	SwinUNETR	AbsLp(p=2)	200	0.3575	0.7163	0.097087	0.0586	0.0457	11.429111	36.2M
18	channels_zt_bhp	SwinUNETR	AbsLp(p=2)	200	0.3605	0.6761	0.084165	0.0574	0.0446	11.190115	16.3M
19	channels_zt_bhp	SwinUNETR	MSE	200	0.3607	0.6707	0.079274	0.0581	0.0456	11.334703	16.3M
20	channels_vt_bhp	SwinUNETR	MSE	200	0.3653	0.7118	0.090665	0.0594	0.0462	11.591043	36.2M
21	channels_vt_rate	SwinUNETR	AbsLp(p=2)	200	0.3670	0.7736	0.101861	0.0841	0.0637	15.169115	36.2M
22	channels_zt_rate	SwinUNETR	MSE	200	0.3685	0.6604	0.076620	0.0872	0.0661	15.608218	16.3M
23	channels_zt_rate	SwinUNETR	AbsLp(p=2)	200	0.3758	0.7278	0.089416	0.0857	0.0647	15.493610	16.3M
24	channels_vt_rate	SwinUNETR	MSE	200	0.3763	0.7227	0.098932	0.0881	0.0663	15.631885	36.2M

Two-Phase Oil-Water

6 results

					Pres (bar)			Sat (frac)
#	Target	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	pressure	SwinUNETR2D	RelLp	200	0.0420	0.0360	6.253993	-	-	-	15.2M
2	combined	SwinUNETR2D	RelLp	200	0.0543	0.0464	7.977582	0.0527	0.0434	0.016646	15.2M
3	pressure	SwinUNETR2D	RelLp	200	0.0554	0.0453	8.037558	-	-	-	6.8M
4	combined	SwinUNETR2D	RelLp	200	0.0566	0.0478	8.175237	0.0547	0.0459	0.016938	6.8M
5	saturation	SwinUNETR2D	RelLp	200	-	-	-	0.0539	0.0431	0.015861	6.8M
6	saturation	SwinUNETR2D	RelLp	200	-	-	-	0.0479	0.0403	0.015003	15.2M

CO2 Radial

3 results

					Pres (bar)			Sat (frac)
#	Target	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	pressure	SwinUNETR2D	RelLp	500	0.0050	0.0013	0.307502	-	-	-	15.2M
2	combined	SwinUNETR2D	RelLp	500	0.0059	0.0017	0.397529	0.0951	8.5406	0.002900	15.2M
3	saturation	SwinUNETR2D	RelLp	500	-	-	-	0.0880	7.7030	0.002304	15.2M

Two-Phase 3D Channels

3 results

					Pres (bar)			Sat (frac)
#	Target	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	pressure	SwinUNETR	RelLp	500	0.0012	0.0009	0.280578	-	-	-	15.7M
2	combined	SwinUNETR	RelLp	500	0.0015	0.0011	0.354836	0.1645	0.0996	0.021278	15.7M
3	saturation	SwinUNETR	RelLp	500	-	-	-	0.1637	0.0890	0.019917	15.7M

Training Curves

SwinUNETR - Flow Arena (channels_nf_bhp)

SwinUNETR - Flow Arena (channels_nf_rate)

SwinUNETR - Flow Arena (channels_vt_bhp)

SwinUNETR - Flow Arena (channels_vt_rate)

SwinUNETR - Flow Arena (channels_zt_bhp)

SwinUNETR - Flow Arena (channels_zt_rate)

SwinUNETR - Flow Arena (geostat_nf_bhp)

SwinUNETR - Flow Arena (geostat_nf_rate)

Architecture

Swin UNETR, introduced by Hatamizadeh et al. (2022), combines the Swin Transformer's shifted window self-attention mechanism with a U-Net-style convolutional decoder. This hybrid architecture aims to capture the best of both worlds: the global context and long-range dependency modeling of transformers with the spatial precision and inductive bias of convolutional decoders.

The encoder is a Swin Transformer that processes the input volume through a hierarchy of shifted window attention blocks. The "shifted window" mechanism restricts self-attention to local windows in one layer and shifts the window partition in the next layer, enabling cross-window information flow without the quadratic cost of full self-attention. The encoder produces multi-scale feature maps that are fed to the convolutional decoder via skip connections.

The decoder follows a standard U-Net transposed-convolution upsampling pathway, progressively combining encoder features with upsampled decoder features through skip connections.

Swin UNETR is the largest model in our benchmark by parameter count (64.1M for 3D), reflecting the overhead of the transformer encoder. It also has the strictest input size requirements: MONAI's implementation requires all spatial dimensions to be divisible by 32 due to the hierarchical window partitioning. This means Swin UNETR cannot be applied to the CO2 Nested dataset (depth=5) without padding or architectural modifications.

Despite its size and constraints, Swin UNETR can excel on datasets where long-range spatial relationships are critical — for example, pressure communication across large reservoir domains.

Model Summary

SwinUNETR

SwinUNETRWrapper(
  (model): SwinUNETR(
    (swinViT): SwinTransformer(
      (patch_embed): PatchEmbed(
        (proj): Conv3d(5, 24, kernel_size=(2, 2, 2), stride=(2, 2, 2))
      )
      (pos_drop): Dropout(p=0.0, inplace=False)
      (layers1): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-1): 2 x SwinTransformerBlock(
              (norm1): LayerNorm((24,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=24, out_features=72, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=24, out_features=24, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((24,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=24, out_features=96, bias=True)
                (linear2): Linear(in_features=96, out_features=24, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=192, out_features=48, bias=False)
            (norm): LayerNorm((192,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
      (layers2): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-1): 2 x SwinTransformerBlock(
              (norm1): LayerNorm((48,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=48, out_features=144, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=48, out_features=48, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((48,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=48, out_features=192, bias=True)
                (linear2): Linear(in_features=192, out_features=48, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=384, out_features=96, bias=False)
            (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
      (layers3): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-5): 6 x SwinTransformerBlock(
              (norm1): LayerNorm((96,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=96, out_features=288, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=96, out_features=96, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((96,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=96, out_features=384, bias=True)
                (linear2): Linear(in_features=384, out_features=96, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=768, out_features=192, bias=False)
            (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
      (layers4): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-1): 2 x SwinTransformerBlock(
              (norm1): LayerNorm((192,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=192, out_features=576, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=192, out_features=192, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((192,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=192, out_features=768, bias=True)
                (linear2): Linear(in_features=768, out_features=192, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=1536, out_features=384, bias=False)
            (norm): LayerNorm((1536,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
    )
    (encoder1): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(5, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(24, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv3d(5, 24, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
        )
        (norm3): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder2): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(24, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(24, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder3): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(48, 48, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(48, 48, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder4): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(96, 96, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(96, 96, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder10): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(384, 384, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(384, 384, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(384, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(384, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder5): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose3d(384, 192, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(384, 192, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(192, 192, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(192, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(192, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv3d(384, 192, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
        )
        (norm3): InstanceNorm3d(192, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder4): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose3d(192, 96, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(192, 96, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(96, 96, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv3d(192, 96, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
        )
        (norm3): InstanceNorm3d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder3): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose3d(96, 48, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(96, 48, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(48, 48, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv3d(96, 48, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
        )
        (norm3): InstanceNorm3d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder2): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose3d(48, 24, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(48, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(24, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv3d(48, 24, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
        )
        (norm3): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder1): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose3d(24, 24, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv3d(48, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv3d(24, 24, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv3d(48, 24, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
        )
        (norm3): InstanceNorm3d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (out): UnetOutBlock(
      (conv): Convolution(
        (conv): Conv3d(24, 40, kernel_size=(1, 1, 1), stride=(1, 1, 1))
      )
    )
  )
)

SwinUNETR2D

SwinUNETR2DWrapper(
  (model): SwinUNETR(
    (swinViT): SwinTransformer(
      (patch_embed): PatchEmbed(
        (proj): Conv2d(84, 24, kernel_size=(2, 2), stride=(2, 2))
      )
      (pos_drop): Dropout(p=0.0, inplace=False)
      (layers1): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-1): 2 x SwinTransformerBlock(
              (norm1): LayerNorm((24,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=24, out_features=72, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=24, out_features=24, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((24,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=24, out_features=96, bias=True)
                (linear2): Linear(in_features=96, out_features=24, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=96, out_features=48, bias=False)
            (norm): LayerNorm((96,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
      (layers2): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-1): 2 x SwinTransformerBlock(
              (norm1): LayerNorm((48,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=48, out_features=144, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=48, out_features=48, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((48,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=48, out_features=192, bias=True)
                (linear2): Linear(in_features=192, out_features=48, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=192, out_features=96, bias=False)
            (norm): LayerNorm((192,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
      (layers3): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-5): 6 x SwinTransformerBlock(
              (norm1): LayerNorm((96,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=96, out_features=288, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=96, out_features=96, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((96,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=96, out_features=384, bias=True)
                (linear2): Linear(in_features=384, out_features=96, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=384, out_features=192, bias=False)
            (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
      (layers4): ModuleList(
        (0): BasicLayer(
          (blocks): ModuleList(
            (0-1): 2 x SwinTransformerBlock(
              (norm1): LayerNorm((192,), eps=1e-05, elementwise_affine=True)
              (attn): WindowAttention(
                (qkv): Linear(in_features=192, out_features=576, bias=True)
                (attn_drop): Dropout(p=0.0, inplace=False)
                (proj): Linear(in_features=192, out_features=192, bias=True)
                (proj_drop): Dropout(p=0.0, inplace=False)
                (softmax): Softmax(dim=-1)
              )
              (drop_path): Identity()
              (norm2): LayerNorm((192,), eps=1e-05, elementwise_affine=True)
              (mlp): MLPBlock(
                (linear1): Linear(in_features=192, out_features=768, bias=True)
                (linear2): Linear(in_features=768, out_features=192, bias=True)
                (fn): GELU(approximate='none')
                (drop1): Dropout(p=0.0, inplace=False)
                (drop2): Dropout(p=0.0, inplace=False)
              )
            )
          )
          (downsample): PatchMerging(
            (reduction): Linear(in_features=768, out_features=384, bias=False)
            (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          )
        )
      )
    )
    (encoder1): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(84, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv2d(84, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
        (norm3): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder2): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder3): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(48, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(48, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder4): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(96, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(96, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (encoder10): UnetrBasicBlock(
      (layer): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(384, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(384, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder5): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose2d(384, 192, kernel_size=(2, 2), stride=(2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(384, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(192, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(192, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv2d(384, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
        (norm3): InstanceNorm2d(192, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder4): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose2d(192, 96, kernel_size=(2, 2), stride=(2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(192, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(96, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv2d(192, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
        (norm3): InstanceNorm2d(96, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder3): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose2d(96, 48, kernel_size=(2, 2), stride=(2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(96, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(48, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv2d(96, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
        (norm3): InstanceNorm2d(48, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder2): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose2d(48, 24, kernel_size=(2, 2), stride=(2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(48, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv2d(48, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
        (norm3): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (decoder1): UnetrUpBlock(
      (transp_conv): Convolution(
        (conv): ConvTranspose2d(24, 24, kernel_size=(2, 2), stride=(2, 2), bias=False)
      )
      (conv_block): UnetResBlock(
        (conv1): Convolution(
          (conv): Conv2d(48, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (conv2): Convolution(
          (conv): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (lrelu): LeakyReLU(negative_slope=0.01, inplace=True)
        (norm1): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (norm2): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
        (conv3): Convolution(
          (conv): Conv2d(48, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
        (norm3): InstanceNorm2d(24, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)
      )
    )
    (out): UnetOutBlock(
      (conv): Convolution(
        (conv): Conv2d(24, 80, kernel_size=(1, 1), stride=(1, 1))
      )
    )
  )
)