class_docstr: Allows segregated segmentation in case the category of an object is known. This is the case in ShapeNet
  for example.
  
      Arguments:
          in_features -- size of the input channel
          cat_to_seg {[type]} -- category to segment maps for example:
              {
                  'Airplane': [0,1,2],
                  'Table': [3,4]
              }
  
      Keyword Arguments:
          dropout_proba  (default: {0.5})
          bn_momentum  -- batch norm momentum (default: {0.1})
      
code: def __init__(
      self, in_features, cat_to_seg, dropout_proba=0.5, bn_momentum=0.1
  ):
      super().__init__()
      self._cat_to_seg = {}
      self._num_categories = len(cat_to_seg)
      self._max_seg_count = 0
      self._max_seg = 0
      self._shifts = torch.zeros(
          (self._num_categories,), dtype=torch.long
      )
      for i, seg in enumerate(cat_to_seg.values()):
          self._max_seg_count = max(self._max_seg_count, len(seg))
          self._max_seg = max(self._max_seg, max(seg))
          self._shifts[i] = min(seg)
          self._cat_to_seg[i] = seg
  
      self.channel_rasing = MLP(
          [in_features, self._num_categories * in_features],
          bn_momentum=bn_momentum,
          bias=False,
      )
      if dropout_proba:
          self.channel_rasing.add_module(
              "Dropout", nn.Dropout(p=dropout_proba)
          )
  
      self.classifier = UnaryConv(
          (self._num_categories, in_features, self._max_seg_count)
      )
      self._bias = Parameter(
          torch.zeros(
              self._max_seg_count,
          )
      )
  
imports: import numpy
  import torch
  from torch import nn
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 1 to create a new MultiHeadClassifier object from inherited base classes: BaseModule with example: obj = MultiHeadClassifier(in_features, cat_to_seg, dropout_proba, bn_momentum)