code: def __init__(self, num_classes):
      super(UNet, self).__init__()
      self.num_classes = num_classes
      self.down1 = UnetEncoder(3, 64)
      self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
      self.down2 = UnetEncoder(64, 128)
      self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
      self.down3 = UnetEncoder(128, 256)
      self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
      self.down4 = UnetEncoder(256, 512)
      self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
      self.center = nn.Sequential(
          nn.Conv2d(512, 1024, kernel_size=3, padding=1),
          nn.BatchNorm2d(1024),
          nn.ReLU(),
          nn.Conv2d(1024, 1024, kernel_size=3, padding=1),
          nn.BatchNorm2d(1024),
          nn.ReLU(),
          nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2),
          nn.BatchNorm2d(512),
          nn.ReLU(),
      )
  
      self.up1 = UnetDecoder(1024, 512, 256)
      self.up2 = UnetDecoder(512, 256, 128)
      self.up3 = UnetDecoder(256, 128, 64)
  
      self.up4 = nn.Sequential(
          nn.Conv2d(128, 64, 3, padding=1),
          # nn.BatchNorm2d(64),
          # nn.ReLU(),
          nn.Conv2d(64, 64, 3, padding=1),
      )
      # nn.BatchNorm2d(64),
      # nn.ReLU())
  
      self.output = nn.Conv2d(
          64, self.num_classes, kernel_size=1, stride=1
      )
      self.final = nn.Sigmoid()
  
      # Initialize weights
      initialize_weights(self)
  
imports: import torch
  import torch.nn
  import torch.nn.init
  import torch.nn.functional
  import numpy
label_desc: Use image node_id 1 to create a new UNet object from inherited base classes: nn with example: obj = UNet(num_classes)