class_docstr: Latent variable modeling of attention head selection
class_docstr_tok: ["Latent","variable","modeling","of","attention","head","selection"]
code: def __init__(
      self,
      num_tasks,
      num_layers,
      total_num_heads,
      num_heads,
      select_strategy="group",
      head_select_temp=5.0,
  ):
      super(AttnHeadSelector, self).__init__()
      self.num_tasks = num_tasks
      self.num_layers = num_layers
      self.total_num_heads = total_num_heads
      self.num_heads = num_heads
      self.select_strategy = select_strategy
      self.temp = head_select_temp
  
      self.head_logits = torch.nn.Parameter(
          torch.Tensor(
              self.num_tasks, self.num_layers, total_num_heads
          ),
          requires_grad=True,
      )
      nn.init.uniform_(
          self.head_logits, a=math.log(0.01), b=math.log(1.0)
      )
  
file_path: fairseq/examples/attention_head_selection/src/modules/attn_head_selector.py
imports: import torch
  import torch.nn
  import math
label_desc: Use image node_id 1 to create a new AttnHeadSelector object from inherited base classes: nn with example: obj = AttnHeadSelector(num_tasks, num_layers, total_num_heads, num_heads, select_strategy, head_select_temp)