code: import torch
  import torch.nn as nn
  import torch.nn.init as init
  
  from .nets.backbone import HourglassBackbone, SuperpointBackbone
  from .nets.junction_decoder import SuperpointDecoder
  from .nets.heatmap_decoder import PixelShuffleDecoder
  from .nets.descriptor_decoder import SuperpointDescriptor
  
  
  def get_model(model_cfg=None, loss_weights=None, mode="train"):
      """ Get model based on the model configuration. """
      # Check dataset config is given
      if model_cfg is None:
          raise ValueError("[Error] The model config is required!")
  
      # List the supported options here
      print("\n\n\t--------Initializing model----------")
      supported_arch = ["simple"]
      if not model_cfg["model_architecture"] in supported_arch:
          raise ValueError(
              "[Error] The model architecture is not in supported arch!")
  
      if model_cfg["model_architecture"] == "simple":
          model = SOLD2Net(model_cfg)
      else:
          raise ValueError(
              "[Error] The model architecture is not in supported arch!")
  
      # Optionally register loss weights to the model
      if mode == "train":
          if loss_weights is not None:
              for param_name, param in loss_weights.items():
                  if isinstance(param, nn.Parameter):
                      print("\t [Debug] Adding %s with value %f to model"
                            % (param_name, param.item()))
                      model.register_parameter(param_name, param)
          else:
              raise ValueError(
                  "[Error] the loss weights can not be None in dynamic weighting mode during training.")
  
      # Display some summary info.
      print("\tModel architecture: %s" % model_cfg["model_architecture"])
      print("\tBackbone: %s" % model_cfg["backbone"])
      print("\tJunction decoder: %s" % model_cfg["junction_decoder"])
      print("\tHeatmap decoder: %s" % model_cfg["heatmap_decoder"])
      print("\t-------------------------------------")
  
      return model
  
  
  class SOLD2Net(nn.Module):
      """ Full network for SOLD². """
      def __init__(self, model_cfg):
          super(SOLD2Net, self).__init__()
          self.name = model_cfg["model_name"]
          self.cfg = model_cfg
  
          # List supported network options
          self.supported_backbone = ["lcnn", "superpoint"]
          self.backbone_net, self.feat_channel = self.get_backbone()
  
          # List supported junction decoder options
          self.supported_junction_decoder = ["superpoint_decoder"]
          self.junction_decoder = self.get_junction_decoder()
  
          # List supported heatmap decoder options
          self.supported_heatmap_decoder = ["pixel_shuffle",
                                            "pixel_shuffle_single"]
          self.heatmap_decoder = self.get_heatmap_decoder()
  
          # List supported descriptor decoder options
          if "descriptor_decoder" in self.cfg:
              self.supported_descriptor_decoder = ["superpoint_descriptor"]
              self.descriptor_decoder = self.get_descriptor_decoder()
  
          # Initialize the model weights
          self.apply(weight_init)
  
      def forward(self, input_images):
          # The backbone
          features = self.backbone_net(input_images)
  
          # junction decoder
          junctions = self.junction_decoder(features)
  
          # heatmap decoder
          heatmaps = self.heatmap_decoder(features)
  
          outputs = {"junctions": junctions, "heatmap": heatmaps}
  
          # Descriptor decoder
          if "descriptor_decoder" in self.cfg:
              outputs["descriptors"] = self.descriptor_decoder(features)
  
          return outputs
  
      def get_backbone(self):
          """ Retrieve the backbone encoder network. """
          if not self.cfg["backbone"] in self.supported_backbone:
              raise ValueError(
                  "[Error] The backbone selection is not supported.")
  
          # lcnn backbone (stacked hourglass)
          if self.cfg["backbone"] == "lcnn":
              backbone_cfg = self.cfg["backbone_cfg"]
              backbone = HourglassBackbone(**backbone_cfg)
              feat_channel = 256
  
          elif self.cfg["backbone"] == "superpoint":
              backbone_cfg = self.cfg["backbone_cfg"]
              backbone = SuperpointBackbone()
              feat_channel = 128
  
          else:
              raise ValueError(
                  "[Error] The backbone selection is not supported.")
  
          return backbone, feat_channel
  
      def get_junction_decoder(self):
          """ Get the junction decoder. """
          if (not self.cfg["junction_decoder"]
              in self.supported_junction_decoder):
              raise ValueError(
                  "[Error] The junction decoder selection is not supported.")
  
          # superpoint decoder
          if self.cfg["junction_decoder"] == "superpoint_decoder":
              decoder = SuperpointDecoder(self.feat_channel,
                                          self.cfg["backbone"])
          else:
              raise ValueError(
                  "[Error] The junction decoder selection is not supported.")
  
          return decoder
  
      def get_heatmap_decoder(self):
          """ Get the heatmap decoder. """
          if not self.cfg["heatmap_decoder"] in self.supported_heatmap_decoder:
              raise ValueError(
                  "[Error] The heatmap decoder selection is not supported.")
  
          # Pixel_shuffle decoder
          if self.cfg["heatmap_decoder"] == "pixel_shuffle":
              if self.cfg["backbone"] == "lcnn":
                  decoder = PixelShuffleDecoder(self.feat_channel,
                                                num_upsample=2)
              elif self.cfg["backbone"] == "superpoint":
                  decoder = PixelShuffleDecoder(self.feat_channel,
                                                num_upsample=3)
              else:
                  raise ValueError("[Error] Unknown backbone option.")
          # Pixel_shuffle decoder with single channel output
          elif self.cfg["heatmap_decoder"] == "pixel_shuffle_single":
              if self.cfg["backbone"] == "lcnn":
                  decoder = PixelShuffleDecoder(
                      self.feat_channel, num_upsample=2, output_channel=1)
              elif self.cfg["backbone"] == "superpoint":
                  decoder = PixelShuffleDecoder(
                      self.feat_channel, num_upsample=3, output_channel=1)
              else:
                  raise ValueError("[Error] Unknown backbone option.")
          else:
              raise ValueError(
                  "[Error] The heatmap decoder selection is not supported.")
  
          return decoder
  
      def get_descriptor_decoder(self):
          """ Get the descriptor decoder. """
          if (not self.cfg["descriptor_decoder"]
              in self.supported_descriptor_decoder):
              raise ValueError(
                  "[Error] The descriptor decoder selection is not supported.")
  
          # SuperPoint descriptor
          if self.cfg["descriptor_decoder"] == "superpoint_descriptor":
              decoder = SuperpointDescriptor(self.feat_channel)
          else:
              raise ValueError(
                  "[Error] The descriptor decoder selection is not supported.")
  
          return decoder
  
  
  def weight_init(m):
      """ Weight initialization function. """
      # Conv2D
      if isinstance(m, nn.Conv2d):
          init.xavier_normal_(m.weight.data)
          if m.bias is not None:
              init.normal_(m.bias.data)
      # Batchnorm
      elif isinstance(m, nn.BatchNorm2d):
          init.normal_(m.weight.data, mean=1, std=0.02)
          init.constant_(m.bias.data, 0)
      # Linear
      elif isinstance(m, nn.Linear):
          init.xavier_normal_(m.weight.data)
          init.normal_(m.bias.data)
      else:
          pass
  
apis: torch.nn.init.normal_ torch.nn.init.constant_ torch.nn.init.xavier_normal_