code: def main(args, print_fn=print):
      print_fn("Experiment arguments: {}".format(args))
  
      if args.random_seed:
          torch.manual_seed(args.random_seed)
      else:
          torch.manual_seed(123)
      # Load dataset
      if args.dataset.startswith("ogbl"):
          graph, split_edge = load_ogb_dataset(args.dataset)
      else:
          raise NotImplementedError
  
      num_nodes = graph.num_nodes()
  
      # set gpu
      if args.gpu_id >= 0 and torch.cuda.is_available():
          device = "cuda:{}".format(args.gpu_id)
      else:
          device = "cpu"
  
      if args.dataset == "ogbl-collab":
          # ogbl-collab dataset is multi-edge graph
          use_coalesce = True
      else:
          use_coalesce = False
  
      # Generate positive and negative edges and corresponding labels
      # Sampling subgraphs and generate node labeling features
      seal_data = SEALData(
          g=graph,
          split_edge=split_edge,
          hop=args.hop,
          neg_samples=args.neg_samples,
          subsample_ratio=args.subsample_ratio,
          use_coalesce=use_coalesce,
          prefix=args.dataset,
          save_dir=args.save_dir,
          num_workers=args.num_workers,
          print_fn=print_fn,
      )
      node_attribute = seal_data.ndata["feat"]
      edge_weight = seal_data.edata["weight"].float()
  
      train_data = seal_data("train")
      val_data = seal_data("valid")
      test_data = seal_data("test")
  
      train_graphs = len(train_data.graph_list)
  
      # Set data loader
  
      train_loader = GraphDataLoader(
          train_data,
          batch_size=args.batch_size,
          num_workers=args.num_workers,
      )
      val_loader = GraphDataLoader(
          val_data,
          batch_size=args.batch_size,
          num_workers=args.num_workers,
      )
      test_loader = GraphDataLoader(
          test_data,
          batch_size=args.batch_size,
          num_workers=args.num_workers,
      )
  
      # set model
      if args.model == "gcn":
          model = GCN(
              num_layers=args.num_layers,
              hidden_units=args.hidden_units,
              gcn_type=args.gcn_type,
              pooling_type=args.pooling,
              node_attributes=node_attribute,
              edge_weights=edge_weight,
              node_embedding=None,
              use_embedding=True,
              num_nodes=num_nodes,
              dropout=args.dropout,
          )
      elif args.model == "dgcnn":
          model = DGCNN(
              num_layers=args.num_layers,
              hidden_units=args.hidden_units,
              k=args.sort_k,
              gcn_type=args.gcn_type,
              node_attributes=node_attribute,
              edge_weights=edge_weight,
              node_embedding=None,
              use_embedding=True,
              num_nodes=num_nodes,
              dropout=args.dropout,
          )
      else:
          raise ValueError("Model error")
  
      model = model.to(device)
      parameters = model.parameters()
      optimizer = torch.optim.Adam(parameters, lr=args.lr)
      loss_fn = BCEWithLogitsLoss()
      print_fn(
          "Total parameters: {}".format(
              sum([p.numel() for p in model.parameters()])
          )
      )
  
      # train and evaluate loop
      summary_val = []
      summary_test = []
      for epoch in range(args.epochs):
          start_time = time.time()
          loss = train(
              model=model,
              dataloader=train_loader,
              loss_fn=loss_fn,
              optimizer=optimizer,
              device=device,
              num_graphs=args.batch_size,
              total_graphs=train_graphs,
          )
          train_time = time.time()
          if epoch % args.eval_steps == 0:
              val_pos_pred, val_neg_pred = evaluate(
                  model=model, dataloader=val_loader, device=device
              )
              test_pos_pred, test_neg_pred = evaluate(
                  model=model, dataloader=test_loader, device=device
              )
  
              val_metric = evaluate_hits(
                  args.dataset, val_pos_pred, val_neg_pred, args.hits_k
              )
              test_metric = evaluate_hits(
                  args.dataset,
                  test_pos_pred,
                  test_neg_pred,
                  args.hits_k,
              )
              evaluate_time = time.time()
              print_fn(
                  "Epoch-{}, train loss: {:.4f}, hits@{}: val-{:.4f}, test-{:.4f}, "
                  "cost time: train-{:.1f}s, total-{:.1f}s".format(
                      epoch,
                      loss,
                      args.hits_k,
                      val_metric,
                      test_metric,
                      train_time - start_time,
                      evaluate_time - start_time,
                  )
              )
              summary_val.append(val_metric)
              summary_test.append(test_metric)
  
      summary_test = np.array(summary_test)
  
      print_fn("Experiment Results:")
      print_fn(
          "Best hits@{}: {:.4f}, epoch: {}".format(
              args.hits_k, np.max(summary_test), np.argmax(summary_test)
          )
      )
  
file_path: dgl/examples/pytorch/seal/main.py
imports: import time
  import numpy
  import torch
  import torch.multiprocessing
  from dgl import EID, NID
  from dgl.dataloading import GraphDataLoader
  from logger import LightLogging
  from model import DGCNN, GCN
  from sampler import SEALData
  from torch.nn import BCEWithLogitsLoss
  from tqdm import tqdm
  from utils import evaluate_hits, load_ogb_dataset, parse_arguments
label_desc: Use image node_id 3 for calling a global function with example usage: main(args, print_fn) without return types