class_docstr: A 2d-Unet implementation with sane defaults.
class_docstr_tok: ["A","2d-Unet","implementation","with","sane","defaults","."]
docstr: :param in_channels: Number of input channels
  :param out_channels: Number of output channels
  :param nf0: Number of features at highest level of U-Net
  :param num_down: Number of downsampling stages.
  :param max_channels: Maximum number of channels (channels multiply by 2 with every downsampling stage)
  :param use_dropout: Whether to use dropout or no.
  :param dropout_prob: Dropout probability if use_dropout=True.
  :param upsampling_mode: Which type of upsampling should be used. See "UpBlock" for documentation.
  :param norm: Which norm to use. If None, no norm is used. Default is Batchnorm with affinity.
  :param outermost_linear: Whether the output layer should be a linear layer or a nonlinear one.
docstr_tok: [":","param","in_channels",":","Number","of","input","channels",":","param","out_channels",":","Number","of","output","channels",":","param","nf0",":","Number","of","features","at","highest","level","of","U-Net",":","param","num_down",":","Number","of","downsampling","stages",".",":","param","max_channels",":","Maximum","number","of","channels","(","channels","multiply","by","2","with","every","downsampling","stage",")",":","param","use_dropout",":","Whether","to","use","dropout","or","no",".",":","param","dropout_prob",":","Dropout","probability","if","use_dropout=True",".",":","param","upsampling_mode",":","Which","type","of","upsampling","should","be","used",".","See","``","UpBlock","''","for","documentation",".",":","param","norm",":","Which","norm","to","use",".","If","None",",","no","norm","is","used",".","Default","is","Batchnorm","with","affinity",".",":","param","outermost_linear",":","Whether","the","output","layer","should","be","a","linear","layer","or","a","nonlinear","one","."]
code: def __init__(
      self,
      in_channels,
      out_channels,
      nf0,
      num_down,
      max_channels,
      use_dropout,
      upsampling_mode="transpose",
      dropout_prob=0.1,
      norm=nn.BatchNorm2d,
      outermost_linear=False,
  ):
      """
      :param in_channels: Number of input channels
      :param out_channels: Number of output channels
      :param nf0: Number of features at highest level of U-Net
      :param num_down: Number of downsampling stages.
      :param max_channels: Maximum number of channels (channels multiply by 2 with every downsampling stage)
      :param use_dropout: Whether to use dropout or no.
      :param dropout_prob: Dropout probability if use_dropout=True.
      :param upsampling_mode: Which type of upsampling should be used. See "UpBlock" for documentation.
      :param norm: Which norm to use. If None, no norm is used. Default is Batchnorm with affinity.
      :param outermost_linear: Whether the output layer should be a linear layer or a nonlinear one.
      """
      super().__init__()
  
      assert (
          num_down > 0
      ), "Need at least one downsampling layer in UNet."
  
      # Define the in block
      self.in_layer = [
          Conv2dSame(
              in_channels,
              nf0,
              kernel_size=3,
              bias=True if norm is None else False,
          )
      ]
      if norm is not None:
          self.in_layer += [norm(nf0, affine=True)]
      self.in_layer += [nn.LeakyReLU(0.2, True)]
  
      if use_dropout:
          self.in_layer += [nn.Dropout2d(dropout_prob)]
      self.in_layer = nn.Sequential(*self.in_layer)
  
      # Define the center UNet block
      self.unet_block = UnetSkipConnectionBlock(
          min(2 ** (num_down - 1) * nf0, max_channels),
          min(2 ** (num_down - 1) * nf0, max_channels),
          use_dropout=use_dropout,
          dropout_prob=dropout_prob,
          norm=None,  # Innermost has no norm (spatial dimension 1)
          upsampling_mode=upsampling_mode,
      )
  
      for i in list(range(0, num_down - 1))[::-1]:
          self.unet_block = UnetSkipConnectionBlock(
              min(2**i * nf0, max_channels),
              min(2 ** (i + 1) * nf0, max_channels),
              use_dropout=use_dropout,
              dropout_prob=dropout_prob,
              submodule=self.unet_block,
              norm=norm,
              upsampling_mode=upsampling_mode,
          )
  
      # Define the out layer. Each unet block concatenates its inputs with its outputs - so the output layer
      # automatically receives the output of the in_layer and the output of the last unet layer.
      self.out_layer = [
          Conv2dSame(
              2 * nf0,
              out_channels,
              kernel_size=3,
              bias=outermost_linear or (norm is None),
          )
      ]
  
      if not outermost_linear:
          if norm is not None:
              self.out_layer += [norm(out_channels, affine=True)]
          self.out_layer += [nn.ReLU(True)]
  
          if use_dropout:
              self.out_layer += [nn.Dropout2d(dropout_prob)]
      self.out_layer = nn.Sequential(*self.out_layer)
  
      self.out_layer_weight = self.out_layer[0].weight
  
imports: import torch
  import torch.nn
  import torchvision.utils
  from torch.nn import functional
label_desc: Use image node_id 1 to create a new Unet object from inherited base classes: nn with example: obj = Unet(in_channels, out_channels, nf0, num_down, max_channels, use_dropout, upsampling_mode, dropout_prob, norm, outermost_linear)