code: def main(args, task):
      warnings.filterwarnings("ignore")
      multilabel_data = {"ppi", "yelp", "amazon"}
      multilabel = args.dataset in multilabel_data
  
      # This flag is excluded for too large dataset, like amazon, the graph of which is too large to be directly
      # shifted to one gpu. So we need to
      # 1. put the whole graph on cpu, and put the subgraphs on gpu in training phase
      # 2. put the model on gpu in training phase, and put the model on cpu in validation/testing phase
      # We need to judge cpu_flag and cuda (below) simultaneously when shift model between cpu and gpu
      if args.dataset in ["amazon"]:
          cpu_flag = True
      else:
          cpu_flag = False
  
      # load and preprocess dataset
      data = load_data(args, multilabel)
      g = data.g
      train_mask = g.ndata["train_mask"]
      val_mask = g.ndata["val_mask"]
      test_mask = g.ndata["test_mask"]
      labels = g.ndata["label"]
  
      train_nid = data.train_nid
  
      in_feats = g.ndata["feat"].shape[1]
      n_classes = data.num_classes
      n_nodes = g.num_nodes()
      n_edges = g.num_edges()
  
      n_train_samples = train_mask.int().sum().item()
      n_val_samples = val_mask.int().sum().item()
      n_test_samples = test_mask.int().sum().item()
  
      print(
          """----Data statistics------'
      #Nodes %d
      #Edges %d
      #Classes/Labels (multi binary labels) %d
      #Train samples %d
      #Val samples %d
      #Test samples %d"""
          % (
              n_nodes,
              n_edges,
              n_classes,
              n_train_samples,
              n_val_samples,
              n_test_samples,
          )
      )
      # load sampler
  
      kwargs = {
          "dn": args.dataset,
          "g": g,
          "train_nid": train_nid,
          "num_workers_sampler": args.num_workers_sampler,
          "num_subg_sampler": args.num_subg_sampler,
          "batch_size_sampler": args.batch_size_sampler,
          "online": args.online,
          "num_subg": args.num_subg,
          "full": args.full,
      }
  
      if args.sampler == "node":
          saint_sampler = SAINTNodeSampler(args.node_budget, **kwargs)
      elif args.sampler == "edge":
          saint_sampler = SAINTEdgeSampler(args.edge_budget, **kwargs)
      elif args.sampler == "rw":
          saint_sampler = SAINTRandomWalkSampler(
              args.num_roots, args.length, **kwargs
          )
      else:
          raise NotImplementedError
      loader = DataLoader(
          saint_sampler,
          collate_fn=saint_sampler.__collate_fn__,
          batch_size=1,
          shuffle=True,
          num_workers=args.num_workers,
          drop_last=False,
      )
      # set device for dataset tensors
      if args.gpu < 0:
          cuda = False
      else:
          cuda = True
          torch.cuda.set_device(args.gpu)
          val_mask = val_mask.cuda()
          test_mask = test_mask.cuda()
          if not cpu_flag:
              g = g.to("cuda:{}".format(args.gpu))
  
      print("labels shape:", g.ndata["label"].shape)
      print("features shape:", g.ndata["feat"].shape)
  
      model = GCNNet(
          in_dim=in_feats,
          hid_dim=args.n_hidden,
          out_dim=n_classes,
          arch=args.arch,
          dropout=args.dropout,
          batch_norm=not args.no_batch_norm,
          aggr=args.aggr,
      )
  
      if cuda:
          model.cuda()
  
      # logger and so on
      log_dir = save_log_dir(args)
      logger = Logger(os.path.join(log_dir, "loggings"))
      logger.write(args)
  
      # use optimizer
      optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
  
      # set train_nids to cuda tensor
      if cuda:
          train_nid = torch.from_numpy(train_nid).cuda()
          print(
              "GPU memory allocated before training(MB)",
              torch.cuda.memory_allocated(device=train_nid.device)
              / 1024
              / 1024,
          )
      start_time = time.time()
      best_f1 = -1
  
      for epoch in range(args.n_epochs):
          for j, subg in enumerate(loader):
              if cuda:
                  subg = subg.to(torch.cuda.current_device())
              model.train()
              # forward
              pred = model(subg)
              batch_labels = subg.ndata["label"]
  
              if multilabel:
                  loss = F.binary_cross_entropy_with_logits(
                      pred,
                      batch_labels,
                      reduction="sum",
                      weight=subg.ndata["l_n"].unsqueeze(1),
                  )
              else:
                  loss = F.cross_entropy(
                      pred, batch_labels, reduction="none"
                  )
                  loss = (subg.ndata["l_n"] * loss).sum()
  
              optimizer.zero_grad()
              loss.backward()
              torch.nn.utils.clip_grad_norm(model.parameters(), 5)
              optimizer.step()
  
              if j == len(loader) - 1:
                  model.eval()
                  with torch.no_grad():
                      train_f1_mic, train_f1_mac = calc_f1(
                          batch_labels.cpu().numpy(),
                          pred.cpu().numpy(),
                          multilabel,
                      )
                      print(
                          f"epoch:{epoch + 1}/{args.n_epochs}, Iteration {j + 1}/"
                          f"{len(loader)}:training loss",
                          loss.item(),
                      )
                      print(
                          "Train F1-mic {:.4f}, Train F1-mac {:.4f}".format(
                              train_f1_mic, train_f1_mac
                          )
                      )
          # evaluate
          model.eval()
          if epoch % args.val_every == 0:
              if (
                  cpu_flag and cuda
              ):  # Only when we have shifted model to gpu and we need to shift it back on cpu
                  model = model.to("cpu")
              val_f1_mic, val_f1_mac = evaluate(
                  model, g, labels, val_mask, multilabel
              )
              print(
                  "Val F1-mic {:.4f}, Val F1-mac {:.4f}".format(
                      val_f1_mic, val_f1_mac
                  )
              )
              if val_f1_mic > best_f1:
                  best_f1 = val_f1_mic
                  print("new best val f1:", best_f1)
                  torch.save(
                      model.state_dict(),
                      os.path.join(
                          log_dir, "best_model_{}.pkl".format(task)
                      ),
                  )
              if cpu_flag and cuda:
                  model.cuda()
  
      end_time = time.time()
      print(f"training using time {end_time - start_time}")
  
      # test
      if args.use_val:
          model.load_state_dict(
              torch.load(
                  os.path.join(
                      log_dir, "best_model_{}.pkl".format(task)
                  )
              )
          )
      if cpu_flag and cuda:
          model = model.to("cpu")
      test_f1_mic, test_f1_mac = evaluate(
          model, g, labels, test_mask, multilabel
      )
      print(
          "Test F1-mic {:.4f}, Test F1-mac {:.4f}".format(
              test_f1_mic, test_f1_mac
          )
      )
  
file_path: dgl/examples/pytorch/graphsaint/train_sampling.py
imports: import argparse
  import os
  import time
  import warnings
  import torch
  import torch.nn.functional
  from config import CONFIG
  from modules import GCNNet
  from sampler import SAINTEdgeSampler, SAINTNodeSampler, SAINTRandomWalkSampler
  from torch.utils.data import DataLoader
  from utils import calc_f1, evaluate, load_data, Logger, save_log_dir
label_desc: Use image node_id 1 for calling a global function with example usage: main(args, task) without return types