docstr: Override `training_step` from PredictionAlgorithm class for CACM-specific training loop.
docstr_tok: ["Override","`","training_step","`","from","PredictionAlgorithm","class","for","CACM-specific","training","loop","."]
code: def training_step(self, train_batch, batch_idx):
      """
      Override `training_step` from PredictionAlgorithm class for CACM-specific training loop.
  
      """
  
      self.featurizer = self.model[0]
      self.classifier = self.model[1]
  
      minibatches = train_batch
  
      objective = 0
      correct, total = 0, 0
      penalty_causal, penalty_conf, penalty_ind, penalty_sel = (
          0,
          0,
          0,
          0,
      )
      nmb = len(minibatches)
  
      features = [self.featurizer(xi) for xi, _, _ in minibatches]
      classifs = [self.classifier(fi) for fi in features]
      targets = [yi for _, yi, _ in minibatches]
  
      for i in range(nmb):
          objective += F.cross_entropy(classifs[i], targets[i])
          correct += (
              (torch.argmax(classifs[i], dim=1) == targets[i])
              .float()
              .sum()
              .item()
          )
          total += classifs[i].shape[0]
  
      objective /= nmb
      loss = objective
  
      if self.attr_types != []:
          for attr_type_idx, attr_type in enumerate(self.attr_types):
              attribute_labels = [
                  ai for _, _, ai in minibatches
              ]  # [(batch_size, num_attrs)_1, batch_size, num_attrs)_2, ..., (batch_size, num_attrs)_(num_environments)]
  
              E_eq_A_attr = attr_type_idx in self.E_eq_A
  
              # Acause regularization
              if attr_type == "causal":
                  penalty_causal += (
                      self.CACMRegularizer.conditional_reg(
                          classifs,
                          [
                              a[:, attr_type_idx]
                              for a in attribute_labels
                          ],
                          [targets],
                          nmb,
                          E_eq_A_attr,
                      )
                  )
  
              # Aconf regularization
              elif attr_type == "conf":
                  penalty_conf += (
                      self.CACMRegularizer.unconditional_reg(
                          classifs,
                          [
                              a[:, attr_type_idx]
                              for a in attribute_labels
                          ],
                          nmb,
                          E_eq_A_attr,
                      )
                  )
  
              # Aind regularization
              elif attr_type == "ind":
                  penalty_ind += self.CACMRegularizer.unconditional_reg(
                      classifs,
                      [a[:, attr_type_idx] for a in attribute_labels],
                      nmb,
                      E_eq_A_attr,
                  )
  
              # Asel regularization
              elif attr_type == "sel":
                  penalty_sel += self.CACMRegularizer.conditional_reg(
                      classifs,
                      [a[:, attr_type_idx] for a in attribute_labels],
                      [targets],
                      nmb,
                      E_eq_A_attr,
                  )
  
          if nmb > 1:
              penalty_causal /= nmb * (nmb - 1) / 2
              penalty_conf /= nmb * (nmb - 1) / 2
              penalty_ind /= nmb * (nmb - 1) / 2
              penalty_sel /= nmb * (nmb - 1) / 2
  
          # Compile loss
          loss += self.lambda_causal * penalty_causal
          loss += self.lambda_conf * penalty_conf
          loss += self.lambda_ind * penalty_ind
          loss += self.lambda_sel * penalty_sel
  
          if torch.is_tensor(penalty_causal):
              penalty_causal = penalty_causal.item()
              self.log(
                  "penalty_causal",
                  penalty_causal,
                  on_step=False,
                  on_epoch=True,
                  prog_bar=True,
              )
          if torch.is_tensor(penalty_conf):
              penalty_conf = penalty_conf.item()
              self.log(
                  "penalty_conf",
                  penalty_conf,
                  on_step=False,
                  on_epoch=True,
                  prog_bar=True,
              )
          if torch.is_tensor(penalty_ind):
              penalty_ind = penalty_ind.item()
              self.log(
                  "penalty_ind",
                  penalty_ind,
                  on_step=False,
                  on_epoch=True,
                  prog_bar=True,
              )
          if torch.is_tensor(penalty_sel):
              penalty_sel = penalty_sel.item()
              self.log(
                  "penalty_sel",
                  penalty_sel,
                  on_step=False,
                  on_epoch=True,
                  prog_bar=True,
              )
  
      elif self.graph is not None:
          pass  # TODO
  
      else:
          raise ValueError("No attribute types or graph provided.")
  
      acc = correct / total
  
      metrics = {"train_acc": acc, "train_loss": loss}
      self.log_dict(
          metrics, on_step=False, on_epoch=True, prog_bar=True
      )
  
      return loss
  
file_path: dowhy/dowhy/causal_prediction/algorithms/cacm.py
imports: import torch
  from torch.nn import functional
  from dowhy.causal_prediction.algorithms.base_algorithm import PredictionAlgorithm
  from dowhy.causal_prediction.algorithms.regularization import Regularizer
label_desc: Use image node_id 2 for calling the CACM obj's underlying member method code with example usage: obj.training_step(train_batch, batch_idx) and returns: loss