code: def main(args):
      if (
          (not args.do_train)
          and (not args.do_valid)
          and (not args.do_test)
          and (not args.evaluate_train)
      ):
          raise ValueError(
              "one of train/val/test mode must be choosed."
          )
  
      if args.init_checkpoint:
          override_config(args)
  
      args.save_path = (
          "log/%s/%s/%s-%s/%s"
          % (
              args.dataset,
              args.model,
              args.hidden_dim,
              args.gamma,
              time.time(),
          )
          if args.save_path == None
          else args.save_path
      )
      writer = SummaryWriter(args.save_path)
  
      # Write logs to checkpoint and console
      set_logger(args)
  
      dataset = LinkPropPredDataset(name=args.dataset)
      split_dict = dataset.get_edge_split()
      nentity = dataset.graph["num_nodes"]
      nrelation = int(max(dataset.graph["edge_reltype"])[0]) + 1
  
      evaluator = Evaluator(name=args.dataset)
  
      args.nentity = nentity
      args.nrelation = nrelation
  
      logging.info("Model: %s" % args.model)
      logging.info("Dataset: %s" % args.dataset)
      logging.info("#entity: %d" % nentity)
      logging.info("#relation: %d" % nrelation)
  
      train_triples = split_dict["train"]
      logging.info("#train: %d" % len(train_triples["head"]))
      valid_triples = split_dict["valid"]
      logging.info("#valid: %d" % len(valid_triples["head"]))
      test_triples = split_dict["test"]
      logging.info("#test: %d" % len(test_triples["head"]))
  
      train_count, train_true_head, train_true_tail = (
          defaultdict(lambda: 4),
          defaultdict(list),
          defaultdict(list),
      )
      for i in tqdm(range(len(train_triples["head"]))):
          head, relation, tail = (
              train_triples["head"][i],
              train_triples["relation"][i],
              train_triples["tail"][i],
          )
          train_count[(head, relation)] += 1
          train_count[(tail, -relation - 1)] += 1
          train_true_head[(relation, tail)].append(head)
          train_true_tail[(head, relation)].append(tail)
  
      kge_model = KGEModel(
          model_name=args.model,
          nentity=nentity,
          nrelation=nrelation,
          hidden_dim=args.hidden_dim,
          gamma=args.gamma,
          double_entity_embedding=args.double_entity_embedding,
          double_relation_embedding=args.double_relation_embedding,
          evaluator=evaluator,
      )
  
      logging.info("Model Parameter Configuration:")
      for name, param in kge_model.named_parameters():
          logging.info(
              "Parameter %s: %s, require_grad = %s"
              % (name, str(param.size()), str(param.requires_grad))
          )
  
      if args.cuda:
          kge_model = kge_model.cuda()
  
      if args.do_train:
          # Set training dataloader iterator
          train_dataloader_head = DataLoader(
              TrainDataset(
                  train_triples,
                  nentity,
                  nrelation,
                  args.negative_sample_size,
                  "head-batch",
                  train_count,
                  train_true_head,
                  train_true_tail,
              ),
              batch_size=args.batch_size,
              shuffle=True,
              num_workers=max(1, args.cpu_num // 2),
              collate_fn=TrainDataset.collate_fn,
          )
  
          train_dataloader_tail = DataLoader(
              TrainDataset(
                  train_triples,
                  nentity,
                  nrelation,
                  args.negative_sample_size,
                  "tail-batch",
                  train_count,
                  train_true_head,
                  train_true_tail,
              ),
              batch_size=args.batch_size,
              shuffle=True,
              num_workers=max(1, args.cpu_num // 2),
              collate_fn=TrainDataset.collate_fn,
          )
  
          train_iterator = BidirectionalOneShotIterator(
              train_dataloader_head, train_dataloader_tail
          )
  
          # Set training configuration
          current_learning_rate = args.learning_rate
          optimizer = torch.optim.Adam(
              filter(lambda p: p.requires_grad, kge_model.parameters()),
              lr=current_learning_rate,
          )
          if args.warm_up_steps:
              warm_up_steps = args.warm_up_steps
          else:
              warm_up_steps = args.max_steps // 2
  
      if args.init_checkpoint:
          # Restore model from checkpoint directory
          logging.info(
              "Loading checkpoint %s..." % args.init_checkpoint
          )
          checkpoint = torch.load(
              os.path.join(args.init_checkpoint, "checkpoint")
          )
          init_step = checkpoint["step"]
          kge_model.load_state_dict(checkpoint["model_state_dict"])
          if args.do_train:
              current_learning_rate = checkpoint[
                  "current_learning_rate"
              ]
              warm_up_steps = checkpoint["warm_up_steps"]
              optimizer.load_state_dict(
                  checkpoint["optimizer_state_dict"]
              )
      else:
          logging.info("Ramdomly Initializing %s Model..." % args.model)
          init_step = 0
  
      step = init_step
  
      logging.info("Start Training...")
      logging.info("init_step = %d" % init_step)
      logging.info("batch_size = %d" % args.batch_size)
      logging.info(
          "negative_adversarial_sampling = %d"
          % args.negative_adversarial_sampling
      )
      logging.info("hidden_dim = %d" % args.hidden_dim)
      logging.info("gamma = %f" % args.gamma)
      logging.info(
          "negative_adversarial_sampling = %s"
          % str(args.negative_adversarial_sampling)
      )
      if args.negative_adversarial_sampling:
          logging.info(
              "adversarial_temperature = %f"
              % args.adversarial_temperature
          )
  
      # Set valid dataloader as it would be evaluated during training
  
      if args.do_train:
          logging.info("learning_rate = %d" % current_learning_rate)
  
          training_logs = []
  
          # Training Loop
          for step in range(init_step, args.max_steps):
              log = kge_model.train_step(
                  kge_model, optimizer, train_iterator, args
              )
              training_logs.append(log)
  
              if step >= warm_up_steps:
                  current_learning_rate = current_learning_rate / 10
                  logging.info(
                      "Change learning_rate to %f at step %d"
                      % (current_learning_rate, step)
                  )
                  optimizer = torch.optim.Adam(
                      filter(
                          lambda p: p.requires_grad,
                          kge_model.parameters(),
                      ),
                      lr=current_learning_rate,
                  )
                  warm_up_steps = warm_up_steps * 3
  
              if (
                  step % args.save_checkpoint_steps == 0 and step > 0
              ):  # ~ 41 seconds/saving
                  save_variable_list = {
                      "step": step,
                      "current_learning_rate": current_learning_rate,
                      "warm_up_steps": warm_up_steps,
                  }
                  save_model(
                      kge_model, optimizer, save_variable_list, args
                  )
  
              if step % args.log_steps == 0:
                  metrics = {}
                  for metric in training_logs[0].keys():
                      metrics[metric] = sum(
                          [log[metric] for log in training_logs]
                      ) / len(training_logs)
                  log_metrics("Train", step, metrics, writer)
                  training_logs = []
  
              if (
                  args.do_valid
                  and step % args.valid_steps == 0
                  and step > 0
              ):
                  logging.info("Evaluating on Valid Dataset...")
                  metrics = kge_model.test_step(
                      kge_model, valid_triples, args
                  )
                  log_metrics("Valid", step, metrics, writer)
  
          save_variable_list = {
              "step": step,
              "current_learning_rate": current_learning_rate,
              "warm_up_steps": warm_up_steps,
          }
          save_model(kge_model, optimizer, save_variable_list, args)
  
      if args.do_valid:
          logging.info("Evaluating on Valid Dataset...")
          metrics = kge_model.test_step(kge_model, valid_triples, args)
          log_metrics("Valid", step, metrics, writer)
  
      if args.do_test:
          logging.info("Evaluating on Test Dataset...")
          metrics = kge_model.test_step(kge_model, test_triples, args)
          log_metrics("Test", step, metrics, writer)
  
      if args.evaluate_train:
          logging.info("Evaluating on Training Dataset...")
          small_train_triples = {}
          indices = np.random.choice(
              len(train_triples["head"]),
              args.ntriples_eval_train,
              replace=False,
          )
          for i in train_triples:
              small_train_triples[i] = train_triples[i][indices]
          metrics = kge_model.test_step(
              kge_model, small_train_triples, args, random_sampling=True
          )
          log_metrics("Train", step, metrics, writer)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import argparse
  import json
  import logging
  import os
  import random
  import numpy
  import torch
  from torch.utils.data import DataLoader
  from model import KGEModel
  from dataloader import TrainDataset
  from dataloader import BidirectionalOneShotIterator
  from ogb.linkproppred import LinkPropPredDataset, Evaluator
  from collections import defaultdict
  from tqdm.auto import tqdm
  import time
  from tensorboardX import SummaryWriter
label_desc: Use image node_id 6 for calling a global function with example usage: main(args) without return types