parameters: (features, labels, mode, params)
docstring: Create the model for estimator api
  
      Args:
          features: if input_layout == 'nhwc', a tensor with shape:
                  [BATCH_SIZE, go.N, go.N, get_features_planes()]
              else, a tensor with shape:
                  [BATCH_SIZE, get_features_planes(), go.N, go.N]
          labels: dict from string to tensor with shape
              'pi_tensor': [BATCH_SIZE, go.N * go.N + 1]
              'value_tensor': [BATCH_SIZE]
          mode: a tf.estimator.ModeKeys (batchnorm params update for TRAIN only)
          params: A dictionary (Typically derived from the FLAGS object.)
      Returns: tf.estimator.EstimatorSpec with props
          mode: same as mode arg
          predictions: dict of tensors
              'policy': [BATCH_SIZE, go.N * go.N + 1]
              'value': [BATCH_SIZE]
          loss: a single value tensor
          train_op: train op
          eval_metric_ops
      return dict of tensors
          logits: [BATCH_SIZE, go.N * go.N + 1]
docstring_summary: Create the model for estimator api
docstring_tokens: Create the model for estimator api
function: def model_fn(features, labels, mode, params):
      """
      Create the model for estimator api
  
      Args:
          features: if input_layout == 'nhwc', a tensor with shape:
                  [BATCH_SIZE, go.N, go.N, get_features_planes()]
              else, a tensor with shape:
                  [BATCH_SIZE, get_features_planes(), go.N, go.N]
          labels: dict from string to tensor with shape
              'pi_tensor': [BATCH_SIZE, go.N * go.N + 1]
              'value_tensor': [BATCH_SIZE]
          mode: a tf.estimator.ModeKeys (batchnorm params update for TRAIN only)
          params: A dictionary (Typically derived from the FLAGS object.)
      Returns: tf.estimator.EstimatorSpec with props
          mode: same as mode arg
          predictions: dict of tensors
              'policy': [BATCH_SIZE, go.N * go.N + 1]
              'value': [BATCH_SIZE]
          loss: a single value tensor
          train_op: train op
          eval_metric_ops
      return dict of tensors
          logits: [BATCH_SIZE, go.N * go.N + 1]
      """
  
      policy_output, value_output, logits = model_inference_fn(
          features, mode == tf.estimator.ModeKeys.TRAIN, params)
  
      # train ops
      policy_cost = tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits_v2(
              logits=logits, labels=tf.stop_gradient(labels['pi_tensor'])))
  
      value_cost = params['value_cost_weight'] * tf.reduce_mean(
          tf.square(value_output - labels['value_tensor']))
  
      reg_vars = [v for v in tf.trainable_variables()
                  if 'bias' not in v.name and 'beta' not in v.name]
      l2_cost = params['l2_strength'] * \
          tf.add_n([tf.nn.l2_loss(v) for v in reg_vars])
  
      combined_cost = policy_cost + value_cost + l2_cost
  
      global_step = tf.train.get_or_create_global_step()
      learning_rate = tf.train.piecewise_constant(
          global_step, params['lr_boundaries'], params['lr_rates'])
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
  
      # Insert quantization ops if requested
      if params['quantize']:
          if mode == tf.estimator.ModeKeys.TRAIN:
              contrib_quantize.create_training_graph(
                  quant_delay=params['quant_delay'])
          else:
              contrib_quantize.create_eval_graph()
  
      optimizer = tf.train.MomentumOptimizer(
          learning_rate, params['sgd_momentum'])
      if params['use_tpu']:
          optimizer = contrib_tpu_python_tpu_tpu_optimizer.CrossShardOptimizer(
              optimizer)
      with tf.control_dependencies(update_ops):
          train_op = optimizer.minimize(combined_cost, global_step=global_step)
  
      # Computations to be executed on CPU, outside of the main TPU queues.
      def eval_metrics_host_call_fn(policy_output, value_output, pi_tensor,
                                    value_tensor, policy_cost, value_cost,
                                    l2_cost, combined_cost, step,
                                    est_mode=tf.estimator.ModeKeys.TRAIN):
          policy_entropy = -tf.reduce_mean(tf.reduce_sum(
              policy_output * tf.log(policy_output), axis=1))
          # pi_tensor is one_hot when generated from sgfs (for supervised learning)
          # and soft-max when using self-play records. argmax normalizes the two.
          policy_target_top_1 = tf.argmax(pi_tensor, axis=1)
  
          policy_output_in_top1 = tf.to_float(
              tf.nn.in_top_k(policy_output, policy_target_top_1, k=1))
          policy_output_in_top3 = tf.to_float(
              tf.nn.in_top_k(policy_output, policy_target_top_1, k=3))
  
          policy_top_1_confidence = tf.reduce_max(policy_output, axis=1)
          policy_target_top_1_confidence = tf.boolean_mask(
              policy_output,
              tf.one_hot(policy_target_top_1, tf.shape(policy_output)[1]))
  
          value_cost_normalized = value_cost / params['value_cost_weight']
          avg_value_observed = tf.reduce_mean(value_tensor)
  
          with tf.variable_scope('metrics'):
              metric_ops = {
                  'policy_cost': tf.metrics.mean(policy_cost),
                  'value_cost': tf.metrics.mean(value_cost),
                  'value_cost_normalized': tf.metrics.mean(value_cost_normalized),
                  'l2_cost': tf.metrics.mean(l2_cost),
                  'policy_entropy': tf.metrics.mean(policy_entropy),
                  'combined_cost': tf.metrics.mean(combined_cost),
                  'avg_value_observed': tf.metrics.mean(avg_value_observed),
                  'policy_accuracy_top_1': tf.metrics.mean(policy_output_in_top1),
                  'policy_accuracy_top_3': tf.metrics.mean(policy_output_in_top3),
                  'policy_top_1_confidence': tf.metrics.mean(policy_top_1_confidence),
                  'policy_target_top_1_confidence': tf.metrics.mean(
                      policy_target_top_1_confidence),
                  'value_confidence': tf.metrics.mean(tf.abs(value_output)),
              }
  
          if est_mode == tf.estimator.ModeKeys.EVAL:
              return metric_ops
  
          # NOTE: global_step is rounded to a multiple of FLAGS.summary_steps.
          eval_step = tf.reduce_min(step)
  
          # Create summary ops so that they show up in SUMMARIES collection
          # That way, they get logged automatically during training
          summary_writer = contrib_summary.create_file_writer(FLAGS.work_dir)
          with summary_writer.as_default(), \
                  contrib_summary.record_summaries_every_n_global_steps(
                      params['summary_steps'], eval_step):
              for metric_name, metric_op in metric_ops.items():
                  contrib_summary.scalar(
                      metric_name, metric_op[1], step=eval_step)
  
          # Reset metrics occasionally so that they are mean of recent batches.
          reset_op = tf.variables_initializer(tf.local_variables('metrics'))
          cond_reset_op = tf.cond(
              tf.equal(eval_step % params['summary_steps'], tf.to_int64(1)),
              lambda: reset_op,
              lambda: tf.no_op())
  
          return contrib_summary.all_summary_ops() + [cond_reset_op]
  
      metric_args = [
          policy_output,
          value_output,
          labels['pi_tensor'],
          labels['value_tensor'],
          tf.reshape(policy_cost, [1]),
          tf.reshape(value_cost, [1]),
          tf.reshape(l2_cost, [1]),
          tf.reshape(combined_cost, [1]),
          tf.reshape(global_step, [1]),
      ]
  
      predictions = {
          'policy_output': policy_output,
          'value_output': value_output,
      }
  
      eval_metrics_only_fn = functools.partial(
          eval_metrics_host_call_fn, est_mode=tf.estimator.ModeKeys.EVAL)
      host_call_fn = functools.partial(
          eval_metrics_host_call_fn, est_mode=tf.estimator.ModeKeys.TRAIN)
  
      tpu_estimator_spec = contrib_tpu_python_tpu_tpu_estimator.TPUEstimatorSpec(
          mode=mode,
          predictions=predictions,
          loss=combined_cost,
          train_op=train_op,
          eval_metrics=(eval_metrics_only_fn, metric_args),
          host_call=(host_call_fn, metric_args)
      )
      if params['use_tpu']:
          return tpu_estimator_spec
      else:
          return tpu_estimator_spec.as_estimator_spec()
function_tokens: def model_fn ( features , labels , mode , params ) : policy_output , value_output , logits = model_inference_fn ( features , mode == tf . estimator . ModeKeys . TRAIN , params ) # train ops policy_cost = tf . reduce_mean ( tf . nn . softmax_cross_entropy_with_logits_v2 ( logits = logits , labels = tf . stop_gradient ( labels [ 'pi_tensor' ] ) ) ) value_cost = params [ 'value_cost_weight' ] * tf . reduce_mean ( tf . square ( value_output - labels [ 'value_tensor' ] ) ) reg_vars = [ v for v in tf . trainable_variables ( ) if 'bias' not in v . name and 'beta' not in v . name ] l2_cost = params [ 'l2_strength' ] * tf . add_n ( [ tf . nn . l2_loss ( v ) for v in reg_vars ] ) combined_cost = policy_cost + value_cost + l2_cost global_step = tf . train . get_or_create_global_step ( ) learning_rate = tf . train . piecewise_constant ( global_step , params [ 'lr_boundaries' ] , params [ 'lr_rates' ] ) update_ops = tf . get_collection ( tf . GraphKeys . UPDATE_OPS ) # Insert quantization ops if requested if params [ 'quantize' ] : if mode == tf . estimator . ModeKeys . TRAIN : contrib_quantize . create_training_graph ( quant_delay = params [ 'quant_delay' ] ) else : contrib_quantize . create_eval_graph ( ) optimizer = tf . train . MomentumOptimizer ( learning_rate , params [ 'sgd_momentum' ] ) if params [ 'use_tpu' ] : optimizer = contrib_tpu_python_tpu_tpu_optimizer . CrossShardOptimizer ( optimizer ) with tf . control_dependencies ( update_ops ) : train_op = optimizer . minimize ( combined_cost , global_step = global_step ) # Computations to be executed on CPU, outside of the main TPU queues. def eval_metrics_host_call_fn ( policy_output , value_output , pi_tensor , value_tensor , policy_cost , value_cost , l2_cost , combined_cost , step , est_mode = tf . estimator . ModeKeys . TRAIN ) : policy_entropy = - tf . reduce_mean ( tf . reduce_sum ( policy_output * tf . log ( policy_output ) , axis = 1 ) ) # pi_tensor is one_hot when generated from sgfs (for supervised learning) # and soft-max when using self-play records. argmax normalizes the two. policy_target_top_1 = tf . argmax ( pi_tensor , axis = 1 ) policy_output_in_top1 = tf . to_float ( tf . nn . in_top_k ( policy_output , policy_target_top_1 , k = 1 ) ) policy_output_in_top3 = tf . to_float ( tf . nn . in_top_k ( policy_output , policy_target_top_1 , k = 3 ) ) policy_top_1_confidence = tf . reduce_max ( policy_output , axis = 1 ) policy_target_top_1_confidence = tf . boolean_mask ( policy_output , tf . one_hot ( policy_target_top_1 , tf . shape ( policy_output ) [ 1 ] ) ) value_cost_normalized = value_cost / params [ 'value_cost_weight' ] avg_value_observed = tf . reduce_mean ( value_tensor ) with tf . variable_scope ( 'metrics' ) : metric_ops = { 'policy_cost' : tf . metrics . mean ( policy_cost ) , 'value_cost' : tf . metrics . mean ( value_cost ) , 'value_cost_normalized' : tf . metrics . mean ( value_cost_normalized ) , 'l2_cost' : tf . metrics . mean ( l2_cost ) , 'policy_entropy' : tf . metrics . mean ( policy_entropy ) , 'combined_cost' : tf . metrics . mean ( combined_cost ) , 'avg_value_observed' : tf . metrics . mean ( avg_value_observed ) , 'policy_accuracy_top_1' : tf . metrics . mean ( policy_output_in_top1 ) , 'policy_accuracy_top_3' : tf . metrics . mean ( policy_output_in_top3 ) , 'policy_top_1_confidence' : tf . metrics . mean ( policy_top_1_confidence ) , 'policy_target_top_1_confidence' : tf . metrics . mean ( policy_target_top_1_confidence ) , 'value_confidence' : tf . metrics . mean ( tf . abs ( value_output ) ) , } if est_mode == tf . estimator . ModeKeys . EVAL : return metric_ops # NOTE: global_step is rounded to a multiple of FLAGS.summary_steps. eval_step = tf . reduce_min ( step ) # Create summary ops so that they show up in SUMMARIES collection # That way, they get logged automatically during training summary_writer = contrib_summary . create_file_writer ( FLAGS . work_dir ) with summary_writer . as_default ( ) , contrib_summary . record_summaries_every_n_global_steps ( params [ 'summary_steps' ] , eval_step ) : for metric_name , metric_op in metric_ops . items ( ) : contrib_summary . scalar ( metric_name , metric_op [ 1 ] , step = eval_step ) # Reset metrics occasionally so that they are mean of recent batches. reset_op = tf . variables_initializer ( tf . local_variables ( 'metrics' ) ) cond_reset_op = tf . cond ( tf . equal ( eval_step % params [ 'summary_steps' ] , tf . to_int64 ( 1 ) ) , lambda : reset_op , lambda : tf . no_op ( ) ) return contrib_summary . all_summary_ops ( ) + [ cond_reset_op ] metric_args = [ policy_output , value_output , labels [ 'pi_tensor' ] , labels [ 'value_tensor' ] , tf . reshape ( policy_cost , [ 1 ] ) , tf . reshape ( value_cost , [ 1 ] ) , tf . reshape ( l2_cost , [ 1 ] ) , tf . reshape ( combined_cost , [ 1 ] ) , tf . reshape ( global_step , [ 1 ] ) , ] predictions = { 'policy_output' : policy_output , 'value_output' : value_output , } eval_metrics_only_fn = functools . partial ( eval_metrics_host_call_fn , est_mode = tf . estimator . ModeKeys . EVAL ) host_call_fn = functools . partial ( eval_metrics_host_call_fn , est_mode = tf . estimator . ModeKeys . TRAIN ) tpu_estimator_spec = contrib_tpu_python_tpu_tpu_estimator . TPUEstimatorSpec ( mode = mode , predictions = predictions , loss = combined_cost , train_op = train_op , eval_metrics = ( eval_metrics_only_fn , metric_args ) , host_call = ( host_call_fn , metric_args ) ) if params [ 'use_tpu' ] : return tpu_estimator_spec else : return tpu_estimator_spec . as_estimator_spec ( )