code: def main(argv: Sequence[str]) -> None:
      del argv
      tf.config.set_soft_device_placement(True)
  
      # Ensure that the first obfuscation in the training obfuscations is 'Clean'.
      if _OBFUSCATIONS_TRAIN.value[0] != data_utils.CLEAN:
          raise ValueError(
              "The first of the training obfuscations must be the (noop)"
              "obfuscation '{}'".format(data_utils.CLEAN)
          )
  
      per_replica_batch_size = _BATCH_SIZE.value
      global_batch_size = per_replica_batch_size * _NUM_CORES.value
      data_config = configs.DatasetConfig(_DATASET.value)
      model_config = configs.ModelConfig(_MODEL_TYPE.value)
  
      train_steps_per_epoch = (
          data_config.train_size // global_batch_size
      )
      eval_steps_per_epoch = int(
          np.ceil(data_config.eval_size / global_batch_size)
      )
  
      resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
          tpu=_TPU.value
      )
      tf.config.experimental_connect_to_cluster(resolver)
      tf.tpu.experimental.initialize_tpu_system(resolver)
      strategy = tf.distribute.experimental.TPUStrategy(resolver)
  
      dataset = data_config.dataset_name_compat
  
      logging.info("Defining dataset.")
      data_train = obfuscations.ObfuscatedEmbeddingDataset(
          data_dir=_DATA_DIR_TRAIN.value,
          embed_dim=model_config.embed_dim,
          split="train",
          batch_size=per_replica_batch_size,
          num_views=len(_OBFUSCATIONS_TRAIN.value),
      )
      train_dataset = strategy.distribute_datasets_from_function(
          data_train.input_fn
      )
  
      data_test = obfuscations.ObfuscatedImageDataset(
          dataset,
          data_dir=_DATA_DIR_EVAL.value,
          obfuscation_list=[data_utils.CLEAN],
          split="test",
          batch_size=per_replica_batch_size,
      )
  
      test_dataset = strategy.distribute_datasets_from_function(
          data_test.input_fn
      )
  
      logging.info("Defined dataset.")
  
      with strategy.scope():
          logging.info("Building model.")
  
          encoder = extended_model.IdentityEmbeddingMapper()
  
          model = extended_model.FeatureExtractor(
              model_link=model_config.model_link,
              bypass_base_model=(_INPUT_FEATURE_NAME.value == "embed"),
              encoder=encoder,
              base_model_trainable=False,
          )
  
          model_clf = extended_model.FeatureExtractorWithClassifier(
              num_classes=data_config.num_classes,
              feature_extractor=model,
              weight_decay=_WEIGHT_DECAY.value,
          )
  
          if _LR_DECAY_TYPE.value == "exponential":
              lr_scheduler = (
                  tf.keras.optimizers.schedules.ExponentialDecay(
                      initial_learning_rate=_BASE_LR.value,
                      decay_steps=train_steps_per_epoch
                      * _LR_DECAY_EPOCHS.value,
                      decay_rate=_LR_DECAY_FACTOR.value,
                      staircase=True,
                  )
              )
          elif _LR_DECAY_TYPE.value == "cosine":
              lr_scheduler = (
                  tf.keras.optimizers.schedules.CosineDecayRestarts(
                      initial_learning_rate=_BASE_LR.value,
                      first_decay_steps=train_steps_per_epoch
                      * _LR_DECAY_EPOCHS.value,
                      alpha=1e-2,
                  )
              )
  
          optimizer = tf.keras.optimizers.SGD(
              learning_rate=lr_scheduler,
              momentum=_MOMENTUM.value,
              nesterov=True,
          )
          logging.info("Built model.")
  
          train_and_eval_model(
              model_clf=model_clf,
              train_dataset=train_dataset,
              train_steps_per_epoch=train_steps_per_epoch,
              test_dataset=test_dataset,
              eval_steps_per_epoch=eval_steps_per_epoch,
              optimizer=optimizer,
              strategy=strategy,
          )
  
file_path: neural-structured-learning/research/meo/linear_finetuning/linear_finetuning.py
imports: import os
  from typing import Sequence
  from absl import app
  from absl import flags
  from absl import logging
  from meo.mlp_baseline import configs
  from meo.mlp_baseline import data_utils
  from meo.mlp_baseline import extended_model
  from meo.mlp_baseline import obfuscations
  import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: main(argv) without return types