class_docstr: This class implements the Graphair model
  
  :param aug_model: The augmentation model g described in the `paper <https://openreview.net/forum?id=1_OGWcP1s9w>`_ used for automated graph augmentations
  :type aug_model: :obj:`torch.nn.Module`
  
  :param f_encoder: The represnetation encoder f described in the `paper <https://openreview.net/forum?id=1_OGWcP1s9w>`_ used for contrastive learning
  :type f_encoder: :obj:`torch.nn.Module`
  
  :param sens_model: The adversary model k described in the `paper <https://openreview.net/forum?id=1_OGWcP1s9w>`_ used for adversarial learning
  :type sens_model: :obj:`torch.nn.Module`
  
  :param classifier_model: The classifier used to predict the sensitive label of nodes on the augmented graph data.
  :type classifier_model: :obj:`torch.nn.Module`
  
  :param lr: Learning rate for aug_model, f_encoder and sens_model. Defaults to 1e-4
  :type lr: float,optional
  
  :param weight_decay: Weight decay for regularization. Defaults to 1e-5
  :type weight_decay: float,optional
  
  :param alpha: The hyperparameter alpha used in the `paper <https://openreview.net/forum?id=1_OGWcP1s9w>`_ to scale adversarial loss component. Defaults to 20.0
  :type alpha: float,optional
  
  :param beta: The hyperparameter beta used in the `paper <https://openreview.net/forum?id=1_OGWcP1s9w>`_ to scale contrastive loss component. Defaults to 0.9
  :type beta: float,optional
  
  :param gamma: The hyperparameter gamma used in the `paper <https://openreview.net/forum?id=1_OGWcP1s9w>`_ to scale reconstruction loss component. Defaults to 0.7
  :type gamma: float,optional
  
  :param lam: The hyperparameter lambda used in the `paper <https://openreview.net/forum?id=1_OGWcP1s9w>`_ to compute reconstruction loss component. Defaults to 1.0
  :type lam: float,optional
  
  :param dataset: The name of the dataset being used. Used only for the model's output path. Defaults to 'POKEC'
  :type dataset: str,optional
  
  :param num_hidden: The input dimension for the MLP networks used in the model. Defaults to 64
  :type num_hidden: int,optional
  
  :param num_proj_hidden: The output dimension for the MLP networks used in the model. Defaults to 64
  :type num_proj_hidden: int,optional
class_docstr_tok: ["This","class","implements","the","Graphair","model",":","param","aug_model",":","The","augmentation","model","g","described","in","the","`","paper","<","https",":","\/\/openreview.net\/forum","?","id=1_OGWcP1s9w",">","`","_","used","for","automated","graph","augmentations",":","type","aug_model",":",":","obj",":","`","torch.nn.Module","`",":","param","f_encoder",":","The","represnetation","encoder","f","described","in","the","`","paper","<","https",":","\/\/openreview.net\/forum","?","id=1_OGWcP1s9w",">","`","_","used","for","contrastive","learning",":","type","f_encoder",":",":","obj",":","`","torch.nn.Module","`",":","param","sens_model",":","The","adversary","model","k","described","in","the","`","paper","<","https",":","\/\/openreview.net\/forum","?","id=1_OGWcP1s9w",">","`","_","used","for","adversarial","learning",":","type","sens_model",":",":","obj",":","`","torch.nn.Module","`",":","param","classifier_model",":","The","classifier","used","to","predict","the","sensitive","label","of","nodes","on","the","augmented","graph","data",".",":","type","classifier_model",":",":","obj",":","`","torch.nn.Module","`",":","param","lr",":","Learning","rate","for","aug_model",",","f_encoder","and","sens_model",".","Defaults","to","1e-4",":","type","lr",":","float",",","optional",":","param","weight_decay",":","Weight","decay","for","regularization",".","Defaults","to","1e-5",":","type","weight_decay",":","float",",","optional",":","param","alpha",":","The","hyperparameter","alpha","used","in","the","`","paper","<","https",":","\/\/openreview.net\/forum","?","id=1_OGWcP1s9w",">","`","_","to","scale","adversarial","loss","component",".","Defaults","to","20.0",":","type","alpha",":","float",",","optional",":","param","beta",":","The","hyperparameter","beta","used","in","the","`","paper","<","https",":","\/\/openreview.net\/forum","?","id=1_OGWcP1s9w",">","`","_","to","scale","contrastive","loss","component",".","Defaults","to","0.9",":","type","beta",":","float",",","optional",":","param","gamma",":","The","hyperparameter","gamma","used","in","the","`","paper","<","https",":","\/\/openreview.net\/forum","?","id=1_OGWcP1s9w",">","`","_","to","scale","reconstruction","loss","component",".","Defaults","to","0.7",":","type","gamma",":","float",",","optional",":","param","lam",":","The","hyperparameter","lambda","used","in","the","`","paper","<","https",":","\/\/openreview.net\/forum","?","id=1_OGWcP1s9w",">","`","_","to","compute","reconstruction","loss","component",".","Defaults","to","1.0",":","type","lam",":","float",",","optional",":","param","dataset",":","The","name","of","the","dataset","being","used",".","Used","only","for","the","model","'s","output","path",".","Defaults","to","'POKEC'",":","type","dataset",":","str",",","optional",":","param","num_hidden",":","The","input","dimension","for","the","MLP","networks","used","in","the","model",".","Defaults","to","64",":","type","num_hidden",":","int",",","optional",":","param","num_proj_hidden",":","The","output","dimension","for","the","MLP","networks","used","in","the","model",".","Defaults","to","64",":","type","num_proj_hidden",":","int",",","optional"]
code: def fit_whole(
      self, epochs, adj, x, sens, idx_sens, warmup=None, adv_epoches=1
  ):
      assert sp.issparse(adj)
      if not isinstance(adj, sp.coo_matrix):
          adj = sp.coo_matrix(adj)
      adj.setdiag(1)
      adj_orig = scipysp_to_pytorchsp(adj).to_dense()
      norm_w = adj_orig.shape[0] ** 2 / float(
          (adj_orig.shape[0] ** 2 - adj_orig.sum()) * 2
      )
      degrees = np.array(adj.sum(1))
      degree_mat_inv_sqrt = sp.diags(np.power(degrees, -0.5).flatten())
      adj_norm = degree_mat_inv_sqrt @ adj @ degree_mat_inv_sqrt
      adj_norm = scipysp_to_pytorchsp(adj_norm)
  
      adj = adj_norm.cuda()
  
      best_contras = float("inf")
  
      if warmup:
          for _ in range(warmup):
              adj_aug, x_aug, adj_logits = self.aug_model(
                  adj, x, adj_orig=adj_orig.cuda()
              )
              edge_loss = norm_w * F.binary_cross_entropy_with_logits(
                  adj_logits, adj_orig.cuda()
              )
  
              feat_loss = self.criterion_recons(x_aug, x)
              recons_loss = edge_loss + self.beta * feat_loss
  
              self.optimizer_aug.zero_grad()
              with torch.autograd.set_detect_anomaly(True):
                  recons_loss.backward(retain_graph=True)
              self.optimizer_aug.step()
  
              print(
                  "edge reconstruction loss: {:.4f}".format(
                      edge_loss.item()
                  ),
                  "feature reconstruction loss: {:.4f}".format(
                      feat_loss.item()
                  ),
              )
  
      for epoch_counter in range(epochs):
          ### generate fair view
          adj_aug, x_aug, adj_logits = self.aug_model(
              adj, x, adj_orig=adj_orig.cuda()
          )
  
          ### extract node representations
          h = self.projection(self.f_encoder(adj, x))
          h_prime = self.projection(self.f_encoder(adj_aug, x_aug))
          # print("encoder done")
  
          ## update sens model
          adj_aug_nograd = adj_aug.detach()
          x_aug_nograd = x_aug.detach()
          if epoch_counter == 0:
              sens_epoches = adv_epoches * 10
          else:
              sens_epoches = adv_epoches
          for _ in range(sens_epoches):
              s_pred, _ = self.sens_model(adj_aug_nograd, x_aug_nograd)
              senloss = self.criterion_sens(
                  s_pred[idx_sens], sens[idx_sens].unsqueeze(1).float()
              )
              self.optimizer_s.zero_grad()
              senloss.backward()
              self.optimizer_s.step()
          s_pred, _ = self.sens_model(adj_aug, x_aug)
          senloss = self.criterion_sens(
              s_pred[idx_sens], sens[idx_sens].unsqueeze(1).float()
          )
  
          ## update aug model
          logits, labels = self.info_nce_loss_2views(
              torch.cat((h, h_prime), dim=0)
          )
          contrastive_loss = self.criterion_cont(logits, labels)
  
          ## update encoder
          edge_loss = norm_w * F.binary_cross_entropy_with_logits(
              adj_logits, adj_orig.cuda()
          )
  
          feat_loss = self.criterion_recons(x_aug, x)
          recons_loss = edge_loss + self.lam * feat_loss
          loss = (
              self.beta * contrastive_loss
              + self.gamma * recons_loss
              - self.alpha * senloss
          )
          self.optimizer.zero_grad()
          loss.backward()
          self.optimizer.step()
  
          print(
              "Epoch: {:04d}".format(epoch_counter + 1),
              "sens loss: {:.4f}".format(senloss.item()),
              "contrastive loss: {:.4f}".format(
                  contrastive_loss.item()
              ),
              "edge reconstruction loss: {:.4f}".format(
                  edge_loss.item()
              ),
              "feature reconstruction loss: {:.4f}".format(
                  feat_loss.item()
              ),
          )
  
      self.save_path = "./checkpoint/graphair_{}_alpha{}_beta{}_gamma{}_lambda{}".format(
          self.dataset, self.alpha, self.beta, self.gamma, self.lam
      )
      torch.save(self.state_dict(), self.save_path)
  
file_path: DIG/dig/fairgraph/method/Graphair/graphair.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  import scipy.sparse
  import numpy
  from dig.fairgraph.utils.utils import scipysp_to_pytorchsp, accuracy, fair_metric
label_desc: Use image node_id 5 for calling the graphair obj's underlying member method code with example usage: obj.fit_whole(epochs, adj, x, sens, idx_sens, warmup, adv_epoches) without return types