code: def main(args):
      # load and preprocess dataset
      data = load_data(args)
      g = data[0]
      features = torch.FloatTensor(g.ndata["feat"])
      labels = torch.LongTensor(g.ndata["label"])
      if hasattr(torch, "BoolTensor"):
          train_mask = torch.BoolTensor(g.ndata["train_mask"])
          val_mask = torch.BoolTensor(g.ndata["val_mask"])
          test_mask = torch.BoolTensor(g.ndata["test_mask"])
      else:
          train_mask = torch.ByteTensor(g.ndata["train_mask"])
          val_mask = torch.ByteTensor(g.ndata["val_mask"])
          test_mask = torch.ByteTensor(g.ndata["test_mask"])
      in_feats = features.shape[1]
      n_classes = data.num_classes
      n_edges = g.num_edges()
  
      if args.gpu < 0:
          cuda = False
      else:
          cuda = True
          torch.cuda.set_device(args.gpu)
          features = features.cuda()
          labels = labels.cuda()
          train_mask = train_mask.cuda()
          val_mask = val_mask.cuda()
          test_mask = test_mask.cuda()
  
      # add self loop
      if args.self_loop:
          g = dgl.remove_self_loop(g)
          g = dgl.add_self_loop(g)
      n_edges = g.num_edges()
  
      if args.gpu >= 0:
          g = g.to(args.gpu)
      # create DGI model
      dgi = DGI(
          g,
          in_feats,
          args.n_hidden,
          args.n_layers,
          nn.PReLU(args.n_hidden),
          args.dropout,
      )
  
      if cuda:
          dgi.cuda()
  
      dgi_optimizer = torch.optim.Adam(
          dgi.parameters(),
          lr=args.dgi_lr,
          weight_decay=args.weight_decay,
      )
  
      # train deep graph infomax
      cnt_wait = 0
      best = 1e9
      best_t = 0
      mean = 0
      for epoch in range(args.n_dgi_epochs):
          dgi.train()
          if epoch >= 3:
              t0 = time.time()
  
          dgi_optimizer.zero_grad()
          loss = dgi(features)
          loss.backward()
          dgi_optimizer.step()
  
          if loss < best:
              best = loss
              best_t = epoch
              cnt_wait = 0
              torch.save(dgi.state_dict(), "best_dgi.pkl")
          else:
              cnt_wait += 1
  
          if cnt_wait == args.patience:
              print("Early stopping!")
              break
  
          if epoch >= 3:
              mean = (mean * (epoch - 3) + (time.time() - t0)) / (
                  epoch - 2
              )
  
              print(
                  "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | "
                  "ETputs(KTEPS) {:.2f}".format(
                      epoch, mean, loss.item(), n_edges / mean / 1000
                  )
              )
  
      # create classifier model
      classifier = Classifier(args.n_hidden, n_classes)
      if cuda:
          classifier.cuda()
  
      classifier_optimizer = torch.optim.Adam(
          classifier.parameters(),
          lr=args.classifier_lr,
          weight_decay=args.weight_decay,
      )
  
      # train classifier
      print("Loading {}th epoch".format(best_t))
      dgi.load_state_dict(torch.load("best_dgi.pkl"))
      embeds = dgi.encoder(features, corrupt=False)
      embeds = embeds.detach()
      mean = 0
      for epoch in range(args.n_classifier_epochs):
          classifier.train()
          if epoch >= 3:
              t0 = time.time()
  
          classifier_optimizer.zero_grad()
          preds = classifier(embeds)
          loss = F.nll_loss(preds[train_mask], labels[train_mask])
          loss.backward()
          classifier_optimizer.step()
  
          if epoch >= 3:
              mean = (mean * (epoch - 3) + (time.time() - t0)) / (
                  epoch - 2
              )
  
              acc = evaluate(classifier, embeds, labels, val_mask)
              print(
                  "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
                  "ETputs(KTEPS) {:.2f}".format(
                      epoch,
                      mean,
                      loss.item(),
                      acc,
                      n_edges / mean / 1000,
                  )
              )
  
      print()
      acc = evaluate(classifier, embeds, labels, test_mask)
      print("Test Accuracy {:.4f}".format(acc))
  
file_path: dgl/examples/pytorch/dgi/train.py
imports: import argparse, time
  import dgl
  import networkx
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from dgi import Classifier, DGI
  from dgl import DGLGraph
  from dgl.data import load_data, register_data_args
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types