class_docstr: ## Margin loss for class existence
  
  A separate margin loss is used for each output capsule and the total loss is the sum of them.
  The length of each output capsule is the probability that class is present in the input.
  
  Loss for each output capsule or class $k$ is,
  $$\mathcal{L}_k = T_k \max(0, m^{+} - \lVert\mathbf{v}_k
  Vert)^2 +
  \lambda (1 - T_k) \max(0, \lVert\mathbf{v}_k
  Vert - m^{-})^2$$
  
  $T_k$ is $1$ if the class $k$ is present and $0$ otherwise.
  The first component of the loss is $0$ when the class is not present,
  and the second component is $0$ if the class is present.
  The $\max(0, x)$ is used to avoid predictions going to extremes.
  $m^{+}$ is set to be $0.9$ and $m^{-}$ to be $0.1$ in the paper.
  
  The $\lambda$ down-weighting is used to stop the length of all capsules from
  falling during the initial phase of training.
class_docstr_tok: ["#","#","Margin","loss","for","class","existence","A","separate","margin","loss","is","used","for","each","output","capsule","and","the","total","loss","is","the","sum","of","them",".","The","length","of","each","output","capsule","is","the","probability","that","class","is","present","in","the","input",".","Loss","for","each","output","capsule","or","class","$","k","$","is",",","$","$","\\mathcal","{","L","}","_k","=","T_k","\\max","(","0",",","m^","{","+","}","-","\\lVert\\mathbf","{","v","}","_k","Vert",")","^2","+","\\lambda","(","1","-","T_k",")","\\max","(","0",",","\\lVert\\mathbf","{","v","}","_k","Vert","-","m^","{","-","}",")","^2","$","$","$","T_k","$","is","$","1","$","if","the","class","$","k","$","is","present","and","$","0","$","otherwise",".","The","first","component","of","the","loss","is","$","0","$","when","the","class","is","not","present",",","and","the","second","component","is","$","0","$","if","the","class","is","present",".","The","$","\\max","(","0",",","x",")","$","is","used","to","avoid","predictions","going","to","extremes",".","$","m^","{","+","}","$","is","set","to","be","$","0.9","$","and","$","m^","{","-","}","$","to","be","$","0.1","$","in","the","paper",".","The","$","\\lambda","$","down-weighting","is","used","to","stop","the","length","of","all","capsules","from","falling","during","the","initial","phase","of","training","."]
docstr: `v`, $\mathbf{v}_j$ are the squashed output capsules.
  This has shape `[batch_size, n_labels, n_features]`; that is, there is a capsule for each label.
  
  `labels` are the labels, and has shape `[batch_size]`.
docstr_tok: ["`","v","`",",","$","\\mathbf","{","v","}","_j","$","are","the","squashed","output","capsules",".","This","has","shape","`","[","batch_size",",","n_labels",",","n_features","]","`",";","that","is",",","there","is","a","capsule","for","each","label",".","`","labels","`","are","the","labels",",","and","has","shape","`","[","batch_size","]","`","."]
code: def forward(self, v: torch.Tensor, labels: torch.Tensor):
      """
      `v`, $\mathbf{v}_j$ are the squashed output capsules.
      This has shape `[batch_size, n_labels, n_features]`; that is, there is a capsule for each label.
  
      `labels` are the labels, and has shape `[batch_size]`.
      """
      # $$\lVert \mathbf{v}_j \rVert$$
      v_norm = torch.sqrt((v**2).sum(dim=-1))
  
      # $$\mathcal{L}$$
      # `labels` is one-hot encoded labels of shape `[batch_size, n_labels]`
      labels = torch.eye(self.n_labels, device=labels.device)[labels]
  
      # $$\mathcal{L}_k = T_k \max(0, m^{+} - \lVert\mathbf{v}_k\rVert)^2 +
      # \lambda (1 - T_k) \max(0, \lVert\mathbf{v}_k\rVert - m^{-})^2$$
      # `loss` has shape `[batch_size, n_labels]`. We have parallelized the computation
      # of $\mathcal{L}_k$ for for all $k$.
      loss = labels * F.relu(
          self.m_positive - v_norm
      ) + self.lambda_ * (1.0 - labels) * F.relu(
          v_norm - self.m_negative
      )
  
      # $$\sum_k \mathcal{L}_k$$
      return loss.sum(dim=-1).mean()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/capsule_networks/__init__.py
imports: import torch.nn
  import torch.nn.functional
  import torch.utils.data
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the MarginLoss obj's underlying member method code with example usage: obj.forward(v, labels) and returns: loss