code: def main(args):
      # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
      # Load dataset
      dataset = dgl.data.FraudDataset(args.dataset, train_size=0.4)
      graph = dataset[0]
      num_classes = dataset.num_classes
  
      # check cuda
      if args.gpu >= 0 and th.cuda.is_available():
          device = "cuda:{}".format(args.gpu)
          args.num_workers = 0
      else:
          device = "cpu"
  
      # retrieve labels of ground truth
      labels = graph.ndata["label"].to(device)
  
      # Extract node features
      feat = graph.ndata["feature"].to(device)
      layers_feat = feat.expand(args.num_layers, -1, -1)
  
      # retrieve masks for train/validation/test
      train_mask = graph.ndata["train_mask"]
      val_mask = graph.ndata["val_mask"]
      test_mask = graph.ndata["test_mask"]
  
      train_idx = (
          th.nonzero(train_mask, as_tuple=False).squeeze(1).to(device)
      )
      val_idx = (
          th.nonzero(val_mask, as_tuple=False).squeeze(1).to(device)
      )
      test_idx = (
          th.nonzero(test_mask, as_tuple=False).squeeze(1).to(device)
      )
  
      # Reinforcement learning module only for positive training nodes
      rl_idx = th.nonzero(
          train_mask.to(device) & labels.bool(), as_tuple=False
      ).squeeze(1)
  
      graph = graph.to(device)
  
      # Step 2: Create model =================================================================== #
      model = CAREGNN(
          in_dim=feat.shape[-1],
          num_classes=num_classes,
          hid_dim=args.hid_dim,
          num_layers=args.num_layers,
          activation=th.tanh,
          step_size=args.step_size,
          edges=graph.canonical_etypes,
      )
  
      model = model.to(device)
  
      # Step 3: Create training components ===================================================== #
      _, cnt = th.unique(labels, return_counts=True)
      loss_fn = th.nn.CrossEntropyLoss(weight=1 / cnt)
      optimizer = optim.Adam(
          model.parameters(), lr=args.lr, weight_decay=args.weight_decay
      )
      if args.early_stop:
          stopper = EarlyStopping(patience=100)
  
      # Step 4: training epochs =============================================================== #
      for epoch in range(args.max_epoch):
          # calculate the distance of each edges and sample based on the distance
          dists = []
          p = []
          for i in range(args.num_layers):
              dist = {}
              graph.ndata["nd"] = th.tanh(
                  model.layers[i].MLP(layers_feat[i])
              )
              for etype in graph.canonical_etypes:
                  graph.apply_edges(_l1_dist, etype=etype)
                  dist[etype] = (
                      graph.edges[etype].data.pop("ed").detach().cpu()
                  )
              dists.append(dist)
              p.append(model.layers[i].p)
          graph.ndata.pop("nd")
          sampler = CARESampler(p, dists, args.num_layers)
  
          # train
          model.train()
          tr_loss = 0
          tr_recall = 0
          tr_auc = 0
          tr_blk = 0
          train_dataloader = dgl.dataloading.DataLoader(
              graph,
              train_idx,
              sampler,
              batch_size=args.batch_size,
              shuffle=True,
              drop_last=False,
              num_workers=args.num_workers,
          )
  
          for input_nodes, output_nodes, blocks in train_dataloader:
              blocks = [b.to(device) for b in blocks]
              train_feature = blocks[0].srcdata["feature"]
              train_label = blocks[-1].dstdata["label"]
              logits_gnn, logits_sim = model(blocks, train_feature)
  
              # compute loss
              blk_loss = loss_fn(
                  logits_gnn, train_label
              ) + args.sim_weight * loss_fn(logits_sim, train_label)
              tr_loss += blk_loss.item()
              tr_recall += recall_score(
                  train_label.cpu(),
                  logits_gnn.argmax(dim=1).detach().cpu(),
              )
              tr_auc += roc_auc_score(
                  train_label.cpu(),
                  softmax(logits_gnn, dim=1)[:, 1].detach().cpu(),
              )
              tr_blk += 1
  
              # backward
              optimizer.zero_grad()
              blk_loss.backward()
              optimizer.step()
  
          # Reinforcement learning module
          model.RLModule(graph, epoch, rl_idx, dists)
  
          # validation
          model.eval()
          val_dataloader = dgl.dataloading.DataLoader(
              graph,
              val_idx,
              sampler,
              batch_size=args.batch_size,
              shuffle=True,
              drop_last=False,
              num_workers=args.num_workers,
          )
  
          val_recall, val_auc, val_loss = evaluate(
              model, loss_fn, val_dataloader, device
          )
  
          # Print out performance
          print(
              "In epoch {}, Train Recall: {:.4f} | Train AUC: {:.4f} | Train Loss: {:.4f}; "
              "Valid Recall: {:.4f} | Valid AUC: {:.4f} | Valid loss: {:.4f}".format(
                  epoch,
                  tr_recall / tr_blk,
                  tr_auc / tr_blk,
                  tr_loss / tr_blk,
                  val_recall,
                  val_auc,
                  val_loss,
              )
          )
  
          if args.early_stop:
              if stopper.step(val_auc, model):
                  break
  
      # Test with mini batch after all epoch
      model.eval()
      if args.early_stop:
          model.load_state_dict(th.load("es_checkpoint.pt"))
      test_dataloader = dgl.dataloading.DataLoader(
          graph,
          test_idx,
          sampler,
          batch_size=args.batch_size,
          shuffle=True,
          drop_last=False,
          num_workers=args.num_workers,
      )
  
      test_recall, test_auc, test_loss = evaluate(
          model, loss_fn, test_dataloader, device
      )
  
      print(
          "Test Recall: {:.4f} | Test AUC: {:.4f} | Test loss: {:.4f}".format(
              test_recall, test_auc, test_loss
          )
      )
  
file_path: dgl/examples/pytorch/caregnn/main_sampling.py
imports: import argparse
  import dgl
  import torch
  import torch.optim
  from model_sampling import _l1_dist, CAREGNN, CARESampler
  from sklearn.metrics import recall_score, roc_auc_score
  from torch.nn.functional import softmax
  from utils import EarlyStopping
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types