code: import os
  import torch
  import torch.nn.functional as F
  import csv
  import glob
  import argparse
  
  from datasets.dataloader import build_loader
  
  
  parser = argparse.ArgumentParser(description="Graph Pooling")
  parser.add_argument('--model', type=str, default="SAGNet", help='model name')
  parser.add_argument('--seed', type=int, default=777, help='seed')
  parser.add_argument('--batch_size', type=int, default=128, help='batch size')
  parser.add_argument('--lr', type=float, default=0.001, help='learning rate')
  parser.add_argument('--weight_decay', type=float, default=0.001, help='weight decay')
  parser.add_argument('--dataset', type=str, default='IMDB-MULTI', help='DD/NCI1/NCI109/Mutagenicity/PROTEINS')
  parser.add_argument('--epochs', type=int, default=1000, help='maximum number of epochs')
  parser.add_argument('--patience', type=int, default=100, help='path to save result')
  
  args = parser.parse_args()
  
  
  def train(args, model, train_loader, val_loader):
      model.train()
      optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
  
      min_loss = 1e10
      max_acc = 0
      patience_cnt = 0
      best_epoch = 0
      val_loss_values = []
      for epoch in range(args.epochs):
          loss_train = 0.0
          loss_dis = 0.0
          correct = 0
          for i, data in enumerate(train_loader):
              optimizer.zero_grad()
              data = data.to(args.device)
              output = model(data)
              cls_loss = F.nll_loss(output, data.y)
              dis_loss = model.compute_disentangle_loss()
              loss = cls_loss + 0.001 * dis_loss 
              loss.backward()
              optimizer.step()
              loss_train += cls_loss.item()
              loss_dis += dis_loss.item() 
              pred = output.max(dim=1)[1]
              correct += pred.eq(data.y).sum().item()
          acc_train = correct / len(train_loader.dataset)
          val_acc, val_loss = test(args, model, val_loader)
          if val_acc > max_acc:
              max_acc = val_acc
          print('Epoch: {:04d}'.format(epoch + 1), 'loss_train: {:.6f}'.format(loss_train), 'loss_dis: {:.6f}'.format(loss_dis),
                'acc_train: {:.6f}'.format(acc_train), 'loss_val: {:.6f}'.format(val_loss),
                'acc_val: {:.6f}'.format(val_acc), 'max_acc: {:.6f}'.format(max_acc))
  
          val_loss_values.append(val_loss)
          if val_loss_values[-1] < min_loss:
              min_loss = val_loss_values[-1]
              best_epoch = epoch
              patience_cnt = 0
              torch.save(model.state_dict(), 'save/' + args.dataset + '/' + str(args.seed) + '.pth')
          else:
              patience_cnt += 1
  
          if patience_cnt == args.patience:
              break
  
      print('Optimization Finished!')
  
  
  def test(args, model, loader):
      model.eval()
      correct = 0.
      loss = 0.
      for data in loader:
          data = data.to(args.device)
          output = model(data)
          pred = output.max(dim=1)[1]
          correct += pred.eq(data.y).sum().item()
          loss += F.nll_loss(output, data.y, reduction='sum').item()
      return correct / len(loader.dataset),loss / len(loader.dataset)
  
  
  def main():
      torch.manual_seed(args.seed)
  
      if torch.cuda.is_available():
          torch.cuda.manual_seed(args.seed)
          args.device = "cuda"
          print('cuda')
      else:
          args.device = "cpu"
  
      train_loader, val_loader, test_loader = build_loader(args)
      
      if args.model == "SAGNet":
          from models.SAGNet import SAGNet, Config
          config = Config()
          model = SAGNet(config, args)
      elif args.model == "GSANet":
          from models.GSANet import GSANet, Config
          config = Config()
          model = GSANet(config, args)
      elif args.model == "HGPSLNet":
          from models.HGPSLNet import HGPSLNet, Config
          config = Config()
          model = HGPSLNet(config, args)
      elif args.model == "ASAPNet":
          from models.ASAPNet import ASAPNet, Config
          config = Config()
          model = ASAPNet(config, args)
      model.to(args.device)
  
      train(args, model, train_loader, val_loader)
  
      model.load_state_dict(torch.load('save/' + args.dataset + '/' + str(args.seed) + '.pth'))
      test_acc, test_loss = test(args, model, test_loader)
      print('Test set results, loss = {:.6f}, accuracy = {:.6f}'.format(test_loss, test_acc))
      
      with open('result.txt', 'a') as f:
          f.write(str(test_acc) + '\n')
  
  
  if __name__ == '__main__':
      main()
  
apis: torch.cuda.manual_seed torch.nn.functional.nll_loss torch.manual_seed torch.cuda.is_available