class_docstr: Generates TaskSpecs for various tasks in Phoenix.
docstr: Creates model spec for all tasks.
docstr_tok: ["Creates","model","spec","for","all","tasks","."]
code: def create_model_spec(
      self,
      features,
      params,
      learning_rate_spec,
      use_tpu,
      trial_mode,
      towers,
      labels,
      mode,
      my_id,
      model_directory,
      predictions_fn,
      optimizer_fn=None,
  ):
      """Creates model spec for all tasks."""
      is_training = mode == tf.estimator.ModeKeys.TRAIN
  
      if not optimizer_fn and is_training:
          final_optimizer_fn = _get_optimizer_fn(
              optimizer=params.optimizer,
              learning_rate=learning_rate_spec["learning_rate"],
              use_tpu=use_tpu,
              exponential_decay_steps=learning_rate_spec.get(
                  "exponential_decay_steps", -1
              ),
              exponential_decay_rate=learning_rate_spec.get(
                  "exponential_decay_rate", -1
              ),
              lr_warmup_steps=self._phoenix_spec.learning_spec.lr_warmup_steps,
          )
      elif is_training:
          final_optimizer_fn = optimizer_fn
          if "params" in inspect.signature(optimizer_fn).parameters:
              final_optimizer_fn = functools.partial(
                  optimizer_fn, params=params
              )
  
      train_op_fn = eval_op_fn
      if is_training:
          train_op_fn = functools.partial(
              _train_op_fn,
              optimizer_fn=final_optimizer_fn,
              use_synchronous_optimizer=self._phoenix_spec.use_synchronous_optimizer,
              l2_regularization=learning_rate_spec.get(
                  "l2_regularization", -1
              ),
              gradient_max_norm=learning_rate_spec.get(
                  "gradient_max_norm", -1
              ),
          )
  
      # One label / One task - Old Phoenix behavior
      if (
          self._phoenix_spec.multi_task_spec is None
          or len(self._phoenix_spec.multi_task_spec) <= 1
      ):
          with tf.compat.v1.variable_scope("Phoenix/Trainer"):
              weights = 1.0
              if (
                  self._phoenix_spec.weight_feature_name
                  and mode != tf.estimator.ModeKeys.PREDICT
              ):
                  weights = features[
                      self._phoenix_spec.weight_feature_name
                  ]
  
              (
                  train_logits_specs,
                  eval_logits_spec,
                  teacher_logits,
              ) = self.get_train_and_eval_logits(
                  towers=towers, trial_mode=trial_mode
              )
  
              final_loss_fn = self._get_loss_fn(
                  original_loss_fn=supply_params_if_needed(
                      self._loss_fn, params
                  ),
                  features=features,
                  mode=mode,
                  my_id=my_id,
                  teacher_logits_spec=teacher_logits,
              )
  
              task_spec = self._create_task_spec(
                  labels=labels,
                  weights=weights,
                  train_logits_specs=train_logits_specs,
                  eval_logits_spec=eval_logits_spec,
                  train_op_fn=train_op_fn,
                  name="Trainer",
                  mode=mode,
                  loss_fn=final_loss_fn,
              )
  
              train_op = None
              if is_training:
                  train_op = _merge_train_op_list(
                      task_spec.train_op_list,
                      self._ensemble_spec.no_train_speedup,
                  )
  
              # create predictions here.
              predictions = predictions_fn(
                  eval_logits_spec.logits,
                  mode=mode,
                  temperature=self._phoenix_spec.temperature,
              )
  
              return ModelSpec(
                  loss=task_spec.loss,
                  train_op=train_op,
                  predictions=predictions,
                  train_hooks=task_spec.train_hooks_list,
                  eval_logits=eval_logits_spec.logits,
              )
  
      # MultiTask. New Phoenix Behavior
      # Details in:
      # go/phoenix-multitask
      num_weights_in_labels = 0
      for task_spec in self._phoenix_spec.multi_task_spec:
          if not task_spec.weight_is_a_feature:
              num_weights_in_labels += 1
  
      # In predict mode we don't have labels
      if labels and not self._phoenix_spec.pass_label_dict_as_is:
          assert (
              len(labels)
              == len(self._phoenix_spec.multi_task_spec)
              + num_weights_in_labels
          )
  
      primary_task = None
      task_spec_list = []
      for task_spec in self._phoenix_spec.multi_task_spec:
          # Build tower on top of searched model for the specific task
          task_towers = collections.defaultdict(list)
          for generator_name, tower_list in towers.items():
              for i, tower_ in enumerate(tower_list):
                  new_tower_name = "{}_{}_{}".format(
                      task_spec.label_name, str(i), generator_name
                  )
                  previous_tower_name = "{}_0_search_generator".format(
                      task_spec.label_name
                  )
                  previous_model_dir = tower_.previous_model_dir
                  is_prior = (
                      generator_name
                      != base_tower_generator.SEARCH_GENERATOR
                  )
                  task_tower = Task.get_task(
                      phoenix_spec=self._phoenix_spec,
                      tower_name=new_tower_name,
                      architecture=np.array(
                          [
                              block_builder.BlockType[block_type]
                              for block_type in task_spec.architecture
                          ]
                      ),
                      is_training=is_training,
                      logits_dimesnion=task_spec.number_of_classes,
                      is_frozen=is_prior,
                      hparams=params,
                      model_directory=model_directory,
                      generator_name=generator_name,
                      previous_tower_name=previous_tower_name,
                      previous_model_dir=previous_model_dir,
                  )
                  base_tensor = tower_.logits_spec.logits
                  if task_spec.apply_activation_on_shared_logits:
                      base_tensor = tf.nn.relu(base_tensor)
                  task_tower(base_tensor, training=is_training)
                  task_towers[generator_name].append(task_tower)
          (
              train_logits_spec,
              eval_logits_spec,
              teacher_logits,
          ) = self.get_train_and_eval_logits(
              task_towers, trial_mode=trial_mode
          )
  
          with tf.compat.v1.variable_scope(
              "Phoenix/trainer_{}".format(task_spec.label_name)
          ):
              if mode != tf.estimator.ModeKeys.PREDICT:
                  task_labels = labels[task_spec.label_name]
                  if self._phoenix_spec.pass_label_dict_as_is:
                      task_labels = labels
              else:
                  task_labels = None
  
              weights = 1.0
              if (
                  task_spec.weight_feature_name
                  and mode != tf.estimator.ModeKeys.PREDICT
              ):
                  if task_spec.weight_is_a_feature:
                      weights = features[task_spec.weight_feature_name]
                  else:
                      weights = labels[task_spec.weight_feature_name]
  
              task_loss_fn = self._get_loss_fn(
                  supply_params_if_needed(
                      extract_task_specific(
                          self._loss_fn, task_spec.label_name
                      ),
                      params,
                  ),
                  features=features,
                  mode=mode,
                  my_id=my_id,
                  teacher_logits_spec=teacher_logits,
              )
              task = self._create_task_spec(
                  labels=task_labels,
                  weights=weights,
                  train_logits_specs=train_logits_spec,
                  eval_logits_spec=eval_logits_spec,
                  train_op_fn=train_op_fn,
                  name=task_spec.label_name,
                  mode=mode,
                  loss_fn=task_loss_fn,
                  add_train_ops=(
                      not self._phoenix_spec.merge_losses_of_multitask
                  ),
              )
              task_spec_list.append(task)
              if (
                  task_spec.label_name
                  == self._phoenix_spec.primary_task_name
              ):
                  primary_task = task
  
      if not self._phoenix_spec.primary_task_name:
          primary_task = task_spec_list[0]
  
      model_spec_predictions = {}
      for task in task_spec_list:
          task_predictions = extract_task_specific(
              predictions_fn, task.name
          )(
              task.logits,
              mode=mode,
              temperature=self._phoenix_spec.temperature,
          )
          for (
              prediction_key,
              prediction_value,
          ) in task_predictions.items():
              prediction_key_name = prediction_key + "/" + task.name
              model_spec_predictions[
                  prediction_key_name
              ] = prediction_value
              if task.name == primary_task.name:
                  model_spec_predictions[
                      prediction_key
                  ] = prediction_value
      logging.info(model_spec_predictions)
  
      train_op_list = []
      train_hooks_list = []
      train_losses = []
      for task in task_spec_list:
          train_op_list.extend(task.train_op_list)
          train_hooks_list.extend(task.train_hooks_list)
          train_losses.extend(task.train_losses)
  
      train_op = None
      if is_training:
          if self._phoenix_spec.merge_losses_of_multitask:
              # In this case, tasks in task_spec_list should not have train_ops nor
              # hooks as in this case we sum up all losses to one and create
              # training ops. Follow the usage of
              # self._phoenix_spec.merge_losses_of_multitask above for when we call
              # self._create_task_spec for more details.
              if train_losses:
                  merged_loss = tf.add_n(
                      train_losses, name="merged_loss"
                  )
                  train_op, train_hooks_list = train_op_fn(merged_loss)
              else:
                  train_op = _merge_train_op_list(
                      [], self._ensemble_spec.no_train_speedup
                  )
          else:
              train_op = _merge_train_op_list(
                  train_op_list, self._ensemble_spec.no_train_speedup
              )
  
      return ModelSpec(
          loss=primary_task.loss,
          train_op=train_op,
          predictions=model_spec_predictions,
          train_hooks=train_hooks_list,
          eval_logits=primary_task.logits,
      )
  
imports: import collections
  import functools
  import inspect
  from absl import logging
  from model_search import block_builder
  from model_search import ensembler
  from model_search.architecture import architecture_utils
  from model_search.architecture import tower
  from model_search.ensembler import EnsembleLogits
  from model_search.generators import base_tower_generator
  from model_search.generators import trial_utils
  from model_search.meta import distillation
  import numpy
  import tensorflow.compat.v2
label_desc: Use image node_id 5 for calling the TaskManager obj's underlying member method code with example usage: obj.create_model_spec(features, params, learning_rate_spec, use_tpu, trial_mode, towers, labels, mode, my_id, model_directory, predictions_fn, optimizer_fn) and returns: ModelSpec, ModelSpec