class_docstr: UNet backbone.
  U-Net: Convolutional Networks for Biomedical Image Segmentation.
  https://arxiv.org/pdf/1505.04597.pdf
  
  Args:
      in_channels (int): Number of input image channels. Default" 3.
      base_channels (int): Number of base channels of each stage.
          The output channels of the first stage. Default: 64.
      num_stages (int): Number of stages in encoder, normally 5. Default: 5.
      strides (Sequence[int 1 | 2]): Strides of each stage in encoder.
          len(strides) is equal to num_stages. Normally the stride of the
          first stage in encoder is 1. If strides[i]=2, it uses stride
          convolution to downsample in the correspondence encoder stage.
          Default: (1, 1, 1, 1, 1).
      enc_num_convs (Sequence[int]): Number of convolutional layers in the
          convolution block of the correspondence encoder stage.
          Default: (2, 2, 2, 2, 2).
      dec_num_convs (Sequence[int]): Number of convolutional layers in the
          convolution block of the correspondence decoder stage.
          Default: (2, 2, 2, 2).
      downsamples (Sequence[int]): Whether use MaxPool to downsample the
          feature map after the first stage of encoder
          (stages: [1, num_stages)). If the correspondence encoder stage use
          stride convolution (strides[i]=2), it will never use MaxPool to
          downsample, even downsamples[i-1]=True.
          Default: (True, True, True, True).
      enc_dilations (Sequence[int]): Dilation rate of each stage in encoder.
          Default: (1, 1, 1, 1, 1).
      dec_dilations (Sequence[int]): Dilation rate of each stage in decoder.
          Default: (1, 1, 1, 1).
      with_cp (bool): Use checkpoint or not. Using checkpoint will save some
          memory while slowing down the training speed. Default: False.
      conv_cfg (dict | None): Config dict for convolution layer.
          Default: None.
      norm_cfg (dict | None): Config dict for normalization layer.
          Default: dict(type='BN').
      act_cfg (dict | None): Config dict for activation layer in ConvModule.
          Default: dict(type='ReLU').
      upsample_cfg (dict): The upsample config of the upsample module in
          decoder. Default: dict(type='InterpConv').
      norm_eval (bool): Whether to set norm layers to eval mode, namely,
          freeze running stats (mean and var). Note: Effect on Batch Norm
          and its variants only. Default: False.
      dcn (bool): Use deformable convolution in convolutional layer or not.
          Default: None.
      plugins (dict): plugins for convolutional layers. Default: None.
  
  Notice:
      The input image size should be divisible by the whole downsample rate
      of the encoder. More detail of the whole downsample rate can be found
      in UNet._check_input_divisible.
code: def __init__(
      self,
      in_channels=3,
      base_channels=64,
      num_stages=5,
      strides=(1, 1, 1, 1, 1),
      enc_num_convs=(2, 2, 2, 2, 2),
      dec_num_convs=(2, 2, 2, 2),
      downsamples=(True, True, True, True),
      enc_dilations=(1, 1, 1, 1, 1),
      dec_dilations=(1, 1, 1, 1),
      with_cp=False,
      conv_cfg=None,
      norm_cfg=dict(type="BN"),
      act_cfg=dict(type="ReLU"),
      upsample_cfg=dict(type="InterpConv"),
      norm_eval=False,
      dcn=None,
      plugins=None,
      init_cfg=[
          dict(type="Kaiming", layer="Conv2d"),
          dict(
              type="Constant", layer=["_BatchNorm", "GroupNorm"], val=1
          ),
      ],
  ):
      super().__init__(init_cfg=init_cfg)
      assert dcn is None, "Not implemented yet."
      assert plugins is None, "Not implemented yet."
      assert len(strides) == num_stages, (
          "The length of strides should be equal to num_stages, "
          f"while the strides is {strides}, the length of "
          f"strides is {len(strides)}, and the num_stages is "
          f"{num_stages}."
      )
      assert len(enc_num_convs) == num_stages, (
          "The length of enc_num_convs should be equal to num_stages, "
          f"while the enc_num_convs is {enc_num_convs}, the length of "
          f"enc_num_convs is {len(enc_num_convs)}, and the num_stages is "
          f"{num_stages}."
      )
      assert len(dec_num_convs) == (num_stages - 1), (
          "The length of dec_num_convs should be equal to (num_stages-1), "
          f"while the dec_num_convs is {dec_num_convs}, the length of "
          f"dec_num_convs is {len(dec_num_convs)}, and the num_stages is "
          f"{num_stages}."
      )
      assert len(downsamples) == (num_stages - 1), (
          "The length of downsamples should be equal to (num_stages-1), "
          f"while the downsamples is {downsamples}, the length of "
          f"downsamples is {len(downsamples)}, and the num_stages is "
          f"{num_stages}."
      )
      assert len(enc_dilations) == num_stages, (
          "The length of enc_dilations should be equal to num_stages, "
          f"while the enc_dilations is {enc_dilations}, the length of "
          f"enc_dilations is {len(enc_dilations)}, and the num_stages is "
          f"{num_stages}."
      )
      assert len(dec_dilations) == (num_stages - 1), (
          "The length of dec_dilations should be equal to (num_stages-1), "
          f"while the dec_dilations is {dec_dilations}, the length of "
          f"dec_dilations is {len(dec_dilations)}, and the num_stages is "
          f"{num_stages}."
      )
      self.num_stages = num_stages
      self.strides = strides
      self.downsamples = downsamples
      self.norm_eval = norm_eval
      self.base_channels = base_channels
  
      self.encoder = nn.ModuleList()
      self.decoder = nn.ModuleList()
  
      for i in range(num_stages):
          enc_conv_block = []
          if i != 0:
              if strides[i] == 1 and downsamples[i - 1]:
                  enc_conv_block.append(nn.MaxPool2d(kernel_size=2))
              upsample = strides[i] != 1 or downsamples[i - 1]
              self.decoder.append(
                  UpConvBlock(
                      conv_block=BasicConvBlock,
                      in_channels=base_channels * 2**i,
                      skip_channels=base_channels * 2 ** (i - 1),
                      out_channels=base_channels * 2 ** (i - 1),
                      num_convs=dec_num_convs[i - 1],
                      stride=1,
                      dilation=dec_dilations[i - 1],
                      with_cp=with_cp,
                      conv_cfg=conv_cfg,
                      norm_cfg=norm_cfg,
                      act_cfg=act_cfg,
                      upsample_cfg=upsample_cfg if upsample else None,
                      dcn=None,
                      plugins=None,
                  )
              )
  
          enc_conv_block.append(
              BasicConvBlock(
                  in_channels=in_channels,
                  out_channels=base_channels * 2**i,
                  num_convs=enc_num_convs[i],
                  stride=strides[i],
                  dilation=enc_dilations[i],
                  with_cp=with_cp,
                  conv_cfg=conv_cfg,
                  norm_cfg=norm_cfg,
                  act_cfg=act_cfg,
                  dcn=None,
                  plugins=None,
              )
          )
          self.encoder.append(nn.Sequential(*enc_conv_block))
          in_channels = base_channels * 2**i
  
imports: import torch
  import torch.nn
  import torch.utils.checkpoint
  from mmcv.cnn import ConvModule, build_norm_layer
  from mmengine.model import BaseModule
  from mmengine.utils.dl_utils.parrots_wrapper import _BatchNorm
  from mmocr.registry import MODELS
label_desc: Use image node_id 1 to create a new UNet object from inherited base classes: BaseModule with example: obj = UNet(in_channels, base_channels, num_stages, strides, enc_num_convs, dec_num_convs, downsamples, enc_dilations, dec_dilations, with_cp, conv_cfg, norm_cfg, act_cfg, upsample_cfg, norm_eval, dcn, plugins, init_cfg)