class_docstr: Swin UNETR based on: "Hatamizadeh et al.,
  Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images
  <https://arxiv.org/abs/2201.01266>"
docstr: Args:
      img_size: spatial dimension of input image.
          This argument is only used for checking that the input image size is divisible by the patch size.
          The tensor passed to forward() can have a dynamic shape as long as its spatial dimensions are divisible by 2**5.
          It will be removed in an upcoming version.
      in_channels: dimension of input channels.
      out_channels: dimension of output channels.
      feature_size: dimension of network feature size.
      depths: number of layers in each stage.
      num_heads: number of attention heads.
      norm_name: feature normalization type and arguments.
      drop_rate: dropout rate.
      attn_drop_rate: attention dropout rate.
      dropout_path_rate: drop path rate.
      normalize: normalize output intermediate features in each stage.
      use_checkpoint: use gradient checkpointing for reduced memory usage.
      spatial_dims: number of spatial dims.
      downsample: module used for downsampling, available options are `"mergingv2"`, `"merging"` and a
          user-specified `nn.Module` following the API defined in :py:class:`monai.networks.nets.PatchMerging`.
          The default is currently `"merging"` (the original version defined in v0.9.0).
      use_v2: using swinunetr_v2, which adds a residual convolution block at the beggining of each swin stage.
  
  Examples::
  
      # for 3D single channel input with size (96,96,96), 4-channel output and feature size of 48.
      >>> net = SwinUNETR(img_size=(96,96,96), in_channels=1, out_channels=4, feature_size=48)
  
      # for 3D 4-channel input with size (128,128,128), 3-channel output and (2,4,2,2) layers in each stage.
      >>> net = SwinUNETR(img_size=(128,128,128), in_channels=4, out_channels=3, depths=(2,4,2,2))
  
      # for 2D single channel input with size (96,96), 2-channel output and gradient checkpointing.
      >>> net = SwinUNETR(img_size=(96,96), in_channels=3, out_channels=2, use_checkpoint=True, spatial_dims=2)
docstr_tok: ["Args",":","img_size",":","spatial","dimension","of","input","image",".","This","argument","is","only","used","for","checking","that","the","input","image","size","is","divisible","by","the","patch","size",".","The","tensor","passed","to","forward","(",")","can","have","a","dynamic","shape","as","long","as","its","spatial","dimensions","are","divisible","by","2","*","*","5",".","It","will","be","removed","in","an","upcoming","version",".","in_channels",":","dimension","of","input","channels",".","out_channels",":","dimension","of","output","channels",".","feature_size",":","dimension","of","network","feature","size",".","depths",":","number","of","layers","in","each","stage",".","num_heads",":","number","of","attention","heads",".","norm_name",":","feature","normalization","type","and","arguments",".","drop_rate",":","dropout","rate",".","attn_drop_rate",":","attention","dropout","rate",".","dropout_path_rate",":","drop","path","rate",".","normalize",":","normalize","output","intermediate","features","in","each","stage",".","use_checkpoint",":","use","gradient","checkpointing","for","reduced","memory","usage",".","spatial_dims",":","number","of","spatial","dims",".","downsample",":","module","used","for","downsampling",",","available","options","are","`","``","mergingv2","''","`",",","`","``","merging","''","`","and","a","user-specified","`","nn.Module","`","following","the","API","defined","in",":","py",":","class",":","`","monai.networks.nets.PatchMerging","`",".","The","default","is","currently","`","``","merging","''","`","(","the","original","version","defined","in","v0.9.0",")",".","use_v2",":","using","swinunetr_v2",",","which","adds","a","residual","convolution","block","at","the","beggining","of","each","swin","stage",".","Examples",":",":","#","for","3D","single","channel","input","with","size","(","96,96,96",")",",","4-channel","output","and","feature","size","of","48",".",">",">",">","net","=","SwinUNETR","(","img_size=","(","96,96,96",")",",","in_channels=1",",","out_channels=4",",","feature_size=48",")","#","for","3D","4-channel","input","with","size","(","128,128,128",")",",","3-channel","output","and","(","2,4,2,2",")","layers","in","each","stage",".",">",">",">","net","=","SwinUNETR","(","img_size=","(","128,128,128",")",",","in_channels=4",",","out_channels=3",",","depths=","(","2,4,2,2",")",")","#","for","2D","single","channel","input","with","size","(","96,96",")",",","2-channel","output","and","gradient","checkpointing",".",">",">",">","net","=","SwinUNETR","(","img_size=","(","96,96",")",",","in_channels=3",",","out_channels=2",",","use_checkpoint=True",",","spatial_dims=2",")"]
code: def __init__(
      self,
      img_size: Sequence[int] | int,
      in_channels: int,
      out_channels: int,
      depths: Sequence[int] = (2, 2, 2, 2),
      num_heads: Sequence[int] = (3, 6, 12, 24),
      feature_size: int = 24,
      norm_name: tuple | str = "instance",
      drop_rate: float = 0.0,
      attn_drop_rate: float = 0.0,
      dropout_path_rate: float = 0.0,
      normalize: bool = True,
      use_checkpoint: bool = False,
      spatial_dims: int = 3,
      downsample="merging",
      use_v2=False,
  ) -> None:
      """
      Args:
          img_size: spatial dimension of input image.
              This argument is only used for checking that the input image size is divisible by the patch size.
              The tensor passed to forward() can have a dynamic shape as long as its spatial dimensions are divisible by 2**5.
              It will be removed in an upcoming version.
          in_channels: dimension of input channels.
          out_channels: dimension of output channels.
          feature_size: dimension of network feature size.
          depths: number of layers in each stage.
          num_heads: number of attention heads.
          norm_name: feature normalization type and arguments.
          drop_rate: dropout rate.
          attn_drop_rate: attention dropout rate.
          dropout_path_rate: drop path rate.
          normalize: normalize output intermediate features in each stage.
          use_checkpoint: use gradient checkpointing for reduced memory usage.
          spatial_dims: number of spatial dims.
          downsample: module used for downsampling, available options are `"mergingv2"`, `"merging"` and a
              user-specified `nn.Module` following the API defined in :py:class:`monai.networks.nets.PatchMerging`.
              The default is currently `"merging"` (the original version defined in v0.9.0).
          use_v2: using swinunetr_v2, which adds a residual convolution block at the beggining of each swin stage.
  
      Examples::
  
          # for 3D single channel input with size (96,96,96), 4-channel output and feature size of 48.
          >>> net = SwinUNETR(img_size=(96,96,96), in_channels=1, out_channels=4, feature_size=48)
  
          # for 3D 4-channel input with size (128,128,128), 3-channel output and (2,4,2,2) layers in each stage.
          >>> net = SwinUNETR(img_size=(128,128,128), in_channels=4, out_channels=3, depths=(2,4,2,2))
  
          # for 2D single channel input with size (96,96), 2-channel output and gradient checkpointing.
          >>> net = SwinUNETR(img_size=(96,96), in_channels=3, out_channels=2, use_checkpoint=True, spatial_dims=2)
  
      """
  
      super().__init__()
  
      img_size = ensure_tuple_rep(img_size, spatial_dims)
      patch_sizes = ensure_tuple_rep(self.patch_size, spatial_dims)
      window_size = ensure_tuple_rep(7, spatial_dims)
  
      if spatial_dims not in (2, 3):
          raise ValueError("spatial dimension should be 2 or 3.")
  
      self._check_input_size(img_size)
  
      if not (0 <= drop_rate <= 1):
          raise ValueError("dropout rate should be between 0 and 1.")
  
      if not (0 <= attn_drop_rate <= 1):
          raise ValueError(
              "attention dropout rate should be between 0 and 1."
          )
  
      if not (0 <= dropout_path_rate <= 1):
          raise ValueError("drop path rate should be between 0 and 1.")
  
      if feature_size % 12 != 0:
          raise ValueError("feature_size should be divisible by 12.")
  
      self.normalize = normalize
  
      self.swinViT = SwinTransformer(
          in_chans=in_channels,
          embed_dim=feature_size,
          window_size=window_size,
          patch_size=patch_sizes,
          depths=depths,
          num_heads=num_heads,
          mlp_ratio=4.0,
          qkv_bias=True,
          drop_rate=drop_rate,
          attn_drop_rate=attn_drop_rate,
          drop_path_rate=dropout_path_rate,
          norm_layer=nn.LayerNorm,
          use_checkpoint=use_checkpoint,
          spatial_dims=spatial_dims,
          downsample=look_up_option(downsample, MERGING_MODE)
          if isinstance(downsample, str)
          else downsample,
          use_v2=use_v2,
      )
  
      self.encoder1 = UnetrBasicBlock(
          spatial_dims=spatial_dims,
          in_channels=in_channels,
          out_channels=feature_size,
          kernel_size=3,
          stride=1,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.encoder2 = UnetrBasicBlock(
          spatial_dims=spatial_dims,
          in_channels=feature_size,
          out_channels=feature_size,
          kernel_size=3,
          stride=1,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.encoder3 = UnetrBasicBlock(
          spatial_dims=spatial_dims,
          in_channels=2 * feature_size,
          out_channels=2 * feature_size,
          kernel_size=3,
          stride=1,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.encoder4 = UnetrBasicBlock(
          spatial_dims=spatial_dims,
          in_channels=4 * feature_size,
          out_channels=4 * feature_size,
          kernel_size=3,
          stride=1,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.encoder10 = UnetrBasicBlock(
          spatial_dims=spatial_dims,
          in_channels=16 * feature_size,
          out_channels=16 * feature_size,
          kernel_size=3,
          stride=1,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.decoder5 = UnetrUpBlock(
          spatial_dims=spatial_dims,
          in_channels=16 * feature_size,
          out_channels=8 * feature_size,
          kernel_size=3,
          upsample_kernel_size=2,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.decoder4 = UnetrUpBlock(
          spatial_dims=spatial_dims,
          in_channels=feature_size * 8,
          out_channels=feature_size * 4,
          kernel_size=3,
          upsample_kernel_size=2,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.decoder3 = UnetrUpBlock(
          spatial_dims=spatial_dims,
          in_channels=feature_size * 4,
          out_channels=feature_size * 2,
          kernel_size=3,
          upsample_kernel_size=2,
          norm_name=norm_name,
          res_block=True,
      )
      self.decoder2 = UnetrUpBlock(
          spatial_dims=spatial_dims,
          in_channels=feature_size * 2,
          out_channels=feature_size,
          kernel_size=3,
          upsample_kernel_size=2,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.decoder1 = UnetrUpBlock(
          spatial_dims=spatial_dims,
          in_channels=feature_size,
          out_channels=feature_size,
          kernel_size=3,
          upsample_kernel_size=2,
          norm_name=norm_name,
          res_block=True,
      )
  
      self.out = UnetOutBlock(
          spatial_dims=spatial_dims,
          in_channels=feature_size,
          out_channels=out_channels,
      )
  
imports: from __future__ import annotations
  import itertools
  from collections.abc import Sequence
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch.nn import LayerNorm
  from typing_extensions import Final
  from monai.networks.blocks import MLPBlock
  from monai.networks.blocks import PatchEmbed, UnetOutBlock, UnetrBasicBlock, UnetrUpBlock
  from monai.networks.layers import DropPath, trunc_normal_
  from monai.utils import ensure_tuple_rep, look_up_option, optional_import
  from monai.utils.deprecate_utils import deprecated_arg
label_desc: Use image node_id 1 to create a new SwinUNETR object from inherited base classes: nn with example: obj = SwinUNETR(img_size, in_channels, out_channels, depths, num_heads, feature_size, norm_name, drop_rate, attn_drop_rate, dropout_path_rate, normalize, use_checkpoint, spatial_dims, downsample, use_v2)