class_docstr: Specific head for training iBOT (https://arxiv.org/pdf/2111.07832.pdf).
  
  Adapted from the official iBOT code base to fit in VISSL:
  https://github.com/bytedance/ibot/blob/da316d82636a7a7356835ef224b13d5f3ace0489/models/head.py#L145
class_docstr_tok: ["Specific","head","for","training","iBOT","(","https",":","\/\/arxiv.org\/pdf\/2111.07832.pdf",")",".","Adapted","from","the","official","iBOT","code","base","to","fit","in","VISSL",":","https",":","\/\/github.com\/bytedance\/ibot\/blob\/da316d82636a7a7356835ef224b13d5f3ace0489\/models\/head.py","#","L145"]
code: def __init__(
      self,
      model_config: AttrDict,
      in_dim: int,
      out_dim: int,
      patch_out_dim: int = 0,
      norm: Optional[str] = None,
      act: str = "gelu",
      last_norm: Optional[str] = None,
      normalize_last_layer: bool = True,
      num_layers: int = 3,
      hidden_dim: int = 2048,
      bottleneck_dim: int = 256,
      shared_head: bool = False,
      **kwargs,
  ):
      super().__init__()
      assert (
          num_layers >= 1
      ), "iBOT head number of layers should be at least 1"
      patch_out_dim = patch_out_dim or out_dim
  
      # TODO(IBOT) - bug in the official repo? each norm should be a
      #  separate instance (not used several times inside Sequential)
      norm = self._build_norm(norm, hidden_dim)
      act = self._build_act(act)
  
      # Build the MLP
      if num_layers == 1:
          self.mlp = nn.Linear(in_dim, bottleneck_dim)
      else:
          layers = [nn.Linear(in_dim, hidden_dim)]
          if norm is not None:
              layers.append(norm)
          layers.append(act)
          for _ in range(num_layers - 2):
              layers.append(nn.Linear(hidden_dim, hidden_dim))
              if norm is not None:
                  layers.append(norm)
              layers.append(act)
          layers.append(nn.Linear(hidden_dim, bottleneck_dim))
          self.mlp = CustomSequential(*layers)
      self.apply(self._init_weights)
  
      # Build the prototypes for class token
      self.prototypes0 = nn.utils.weight_norm(
          nn.Linear(bottleneck_dim, out_dim, bias=False)
      )
      self.prototypes0.weight_g.data.fill_(1)
      if normalize_last_layer:
          self.prototypes0.weight_g.requires_grad = False
  
      # Build the prototypes for patch tokens
      if not shared_head:
          self.prototypes1 = nn.utils.weight_norm(
              nn.Linear(bottleneck_dim, patch_out_dim, bias=False)
          )
          self.prototypes1.weight_g.data.fill_(1)
          if normalize_last_layer:
              self.prototypes1.weight_g.requires_grad = False
      else:
          self.prototypes1 = self.prototypes0
  
      # Build the last norm for class token and patch token
      self.last_norm = self._build_norm(
          last_norm, out_dim, affine=False, **kwargs
      )
      if not shared_head:
          self.last_norm2 = self._build_norm(
              last_norm, patch_out_dim, affine=False, **kwargs
          )
      else:
          self.last_norm2 = self.last_norm
  
file_path: vissl/vissl/models/heads/ibot_head.py
imports: from typing import Optional
  import torch
  import torch.nn
  from vissl.config import AttrDict
  from vissl.models.heads import register_model_head
  from vissl.models.model_helpers import trunc_normal_
label_desc: Use image node_id 1 to create a new IBOTHead object from inherited base classes: nn with example: obj = IBOTHead(model_config, in_dim, out_dim, patch_out_dim, norm, act, last_norm, normalize_last_layer, num_layers, hidden_dim, bottleneck_dim, shared_head)