code: def main(argv):
      # Check that the correct number of arguments have been provided. The
      # training and test data should contain 'tf.train.Example' objects in the
      # TFRecord format.
      if len(argv) != 3:
          raise app.UsageError(
              "Invalid number of arguments; expected 2, got %d"
              % (len(argv) - 1)
          )
      if not FLAGS.output_dir:
          raise app.UsageError("Must specify the output dir")
  
      hparams = get_hyper_parameters()
      train_data_path = argv[1]
      test_data_path = argv[2]
  
      # Graph regularization configuration.
      graph_reg_config = nsl.configs.make_graph_reg_config(
          max_neighbors=hparams.num_neighbors,
          multiplier=hparams.graph_regularization_multiplier,
          distance_type=hparams.distance_type,
          sum_over_axis=-1,
      )
  
      # Create the base MLP models.
      base_model_tags = ["FUNCTIONAL", "SEQUENTIAL", "SUBCLASS"]
      for base_model_tag in base_model_tags:
          base_model = build_model(base_model_tag, hparams)
          logging.info(
              "\n====== %s BASE MODEL TEST BEGIN ======", base_model_tag
          )
          train_dataset = make_dataset(
              train_data_path, True, False, hparams
          )
          test_dataset = make_dataset(
              test_data_path, False, False, hparams
          )
          train_and_evaluate(
              base_model,
              "Base MLP model",
              None,
              train_dataset,
              test_dataset,
              hparams,
          )
  
          logging.info(
              "\n====== TRAINING WITH GRAPH REGULARIZATION ======\n"
          )
          kbs_server, kbs_address = start_server()
          config = kbs_config(hparams.num_classes)
          client = de_nb_cache.DynamicEmbeddingNeighborCache(
              "id", config, kbs_address, timeout_ms=10 * 1000
          )
          graph_reg_model = (
              graph_regularization.GraphRegularizationWithCaching(
                  base_model, graph_reg_config, client
              )
          )
  
          train_dataset = make_dataset(
              train_data_path, True, True, hparams
          )
          test_dataset = make_dataset(
              test_data_path, False, False, hparams
          )
          graph_model_path = os.path.join(
              FLAGS.output_dir, "model_" + base_model_tag
          )
  
          # Start a background thread to update emebddings.
          stop_event = threading.Event()
          base_inference_model = build_model(base_model_tag, hparams)
          graph_inference_model = (
              graph_regularization.GraphRegularizationWithCaching(
                  base_inference_model, graph_reg_config, client
              )
          )
          update_thread = threading.Thread(
              target=update_embeddings,
              args=(
                  stop_event,
                  graph_model_path,
                  graph_inference_model,
                  client,
                  train_dataset,
              ),
          )
          update_thread.daemon = True
          update_thread.start()
  
          # Train graph-regularized model using historical neighbor embeddings.
          train_and_evaluate(
              graph_reg_model,
              "MLP + graph regularization",
              graph_model_path,
              train_dataset,
              test_dataset,
              hparams,
          )
          stop_event.set()
          kbs_server.Terminate()
          update_thread.join()
  
          logging.info(
              "\n====== %s BASE MODEL TEST END ======", base_model_tag
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import threading
  from absl import app
  from absl import flags
  from absl import logging
  import attr
  import neural_structured_learning
  from research.carls import dynamic_embedding_config_pb2
  from research.carls import dynamic_embedding_neighbor_cache
  from research.carls import graph_regularization
  from research.carls import kbs_server_helper_pybind
  import tensorflow
  from google.protobuf import text_format
label_desc: Use image node_id 13 for calling a global function with example usage: main(argv) without return types