class_docstr: Specific head for distilling iBOT (https://arxiv.org/pdf/2111.07832.pdf).
class_docstr_tok: ["Specific","head","for","distilling","iBOT","(","https",":","\/\/arxiv.org\/pdf\/2111.07832.pdf",")","."]
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,
      num_layers: int = 3,
      hidden_dim: int = 2048,
      shared_head: bool = False,
      **kwargs,
  ):
      super().__init__()
      assert (
          num_layers >= 2
      ), "iBOT distillation head number of layers should be at least 2"
      patch_out_dim = patch_out_dim or out_dim
  
      norm = self._build_norm(norm, hidden_dim)
      act = self._build_act(act)
  
      # Build the MLP
      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)
      self.mlp = CustomSequential(*layers)
  
      # Build the projections
      self.to_class_out = nn.Linear(hidden_dim, out_dim)
      if not shared_head:
          self.to_patch_out = nn.Linear(hidden_dim, patch_out_dim)
      else:
          self.to_patch_out = self.to_class_out
  
      # Initialize the weights
      self.apply(self._init_weights)
  
      # 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_distillation_head.py
imports: from typing import Optional
  import torch.nn
  from vissl.config import AttrDict
  from vissl.models.heads import register_model_head
  from vissl.models.heads.ibot_head import CustomSequential
  from vissl.models.model_helpers import trunc_normal_
label_desc: Use image node_id 1 to create a new IBOTDistillationHead object from inherited base classes: nn with example: obj = IBOTDistillationHead(model_config, in_dim, out_dim, patch_out_dim, norm, act, last_norm, num_layers, hidden_dim, shared_head)