parameters: (num_gpus, variable_strategy, num_workers)
return_statement: return _mlp_model_fn
docstring: Returns a function that will build the mlp model.
docstring_summary: Returns a function that will build the mlp model.
docstring_tokens: Returns a function that will build the mlp model .
function: def get_model_fn(num_gpus, variable_strategy, num_workers):
      """Returns a function that will build the mlp model."""
  
      def _mlp_model_fn(features, labels, mode, params):
          """Mlp model body.
  
          Support single host, one or more GPU training. Parameter distribution can
          be either one of the following scheme.
          1. CPU is the parameter server and manages gradient updates.
          2. Parameters are distributed evenly across all GPUs, and the first GPU
               manages gradient updates.
  
          Args:
              features: a list of tensors, one for each tower
              labels: a list of tensors, one for each tower
              mode: ModeKeys.TRAIN or EVAL
              params: Hyperparameters suitable for tuning
          Returns:
              A EstimatorSpec object.
          """
          is_training = (mode == tf.estimator.ModeKeys.TRAIN)
          weight_decay = params.weight_decay
          momentum = params.momentum
  
          tower_features = features
          tower_labels = labels
          tower_losses = []
          tower_gradvars = []
          tower_preds = []
  
          # channels first (NCHW) is normally optimal on GPU and channels last (NHWC)
          # on CPU. The exception is Intel MKL on CPU which is optimal with
          # channels_last.
          data_format = params.data_format
          if not data_format:
              if num_gpus == 0:
                  data_format = 'channels_last'
              else:
                  data_format = 'channels_first'
  
          if num_gpus == 0:
              num_devices = 1
              device_type = 'cpu'
          else:
              num_devices = num_gpus
              device_type = 'gpu'
  
          for i in range(num_devices):
              worker_device = '/{}:{}'.format(device_type, i)
              if variable_strategy == 'CPU':
                  device_setter = mlp_utils.local_device_setter(
                      worker_device=worker_device)
              elif variable_strategy == 'GPU':
                  device_setter = mlp_utils.local_device_setter(
                      ps_device_type='gpu',
                      worker_device=worker_device,
                      ps_strategy=tf.contrib.training.
                      GreedyLoadBalancingStrategy(
                          num_gpus, tf.contrib.training.byte_size_load_fn))
              with tf.variable_scope('mlp', reuse=bool(i != 0)):
                  with tf.name_scope('tower_%d' % i) as name_scope:
                      with tf.device(device_setter):
                          loss, gradvars, preds = _tower_fn(
                              is_training, weight_decay, tower_features[i],
                              tower_labels[i], data_format,
                              params.batch_norm_decay, params.batch_norm_epsilon)
                          tower_losses.append(loss)
                          tower_gradvars.append(gradvars)
                          tower_preds.append(preds)
                          if i == 0:
                              # Only trigger batch_norm moving mean and variance update from
                              # the 1st tower. Ideally, we should grab the updates from all
                              # towers but these stats accumulate extremely fast so we can
                              # ignore the other stats from the other towers without
                              # significant detriment.
                              update_ops = tf.get_collection(
                                  tf.GraphKeys.UPDATE_OPS, name_scope)
  
          # Now compute global loss and gradients.
          gradvars = []
          with tf.name_scope('gradient_averaging'):
              all_grads = {}
              for grad, var in itertools.chain(*tower_gradvars):
                  if grad is not None:
                      all_grads.setdefault(var, []).append(grad)
              for var, grads in six.iteritems(all_grads):
                  # Average gradients on the same device as the variables
                  # to which they apply.
                  with tf.device(var.device):
                      if len(grads) == 1:
                          avg_grad = grads[0]
                      else:
                          avg_grad = tf.multiply(
                              tf.add_n(grads), 1. / len(grads))
                  gradvars.append((avg_grad, var))
  
          # Device that runs the ops to apply global gradient updates.
          consolidation_device = '/gpu:0' if variable_strategy == 'GPU' else '/cpu:0'
          with tf.device(consolidation_device):
              num_batches_per_epoch = mlp_data.MlpDataSet.num_examples_per_epoch(
                  'train') // (params.train_batch_size * num_workers)
              boundaries = [
                  num_batches_per_epoch * x
                  for x in np.array([20, 50, 80], dtype=np.int64)
              ]
              staged_lr = [
                  params.learning_rate * x for x in [1, 0.1, 0.01, 0.002]
              ]
  
              learning_rate = tf.train.piecewise_constant(
                  tf.train.get_global_step(), boundaries, staged_lr)
  
              loss = tf.reduce_mean(tower_losses, name='loss')
  
              examples_sec_hook = mlp_utils.ExamplesPerSecondHook(
                  params.train_batch_size, every_n_steps=10)
  
              tensors_to_log = {'learning_rate': learning_rate, 'loss': loss}
  
              logging_hook = tf.train.LoggingTensorHook(
                  tensors=tensors_to_log, every_n_iter=100)
  
              train_hooks = [logging_hook, examples_sec_hook]
  
              optimizer = tf.train.MomentumOptimizer(
                  learning_rate=learning_rate, momentum=momentum)
  
              if params.sync:
                  optimizer = tf.train.SyncReplicasOptimizer(
                      optimizer, replicas_to_aggregate=num_workers)
                  sync_replicas_hook = optimizer.make_session_run_hook(
                      params.is_chief)
                  train_hooks.append(sync_replicas_hook)
  
              # Create single grouped train op
              train_op = [
                  optimizer.apply_gradients(
                      gradvars, global_step=tf.train.get_global_step())
              ]
              train_op.extend(update_ops)
              train_op = tf.group(*train_op)
  
              predictions = {
                  'classes':
                  tf.concat([p['classes'] for p in tower_preds], axis=0),
                  'probabilities':
                  tf.concat([p['probabilities'] for p in tower_preds], axis=0)
              }
              stacked_labels = tf.concat(labels, axis=0)
              metrics = {
                  'accuracy':
                  tf.metrics.accuracy(stacked_labels, predictions['classes'])
              }
  
          return tf.estimator.EstimatorSpec(
              mode=mode,
              predictions=predictions,
              loss=loss,
              train_op=train_op,
              training_hooks=train_hooks,
              eval_metric_ops=metrics)
  
      return _mlp_model_fn
function_tokens: def get_model_fn ( num_gpus , variable_strategy , num_workers ) : def _mlp_model_fn ( features , labels , mode , params ) : """Mlp model body.
  
          Support single host, one or more GPU training. Parameter distribution can
          be either one of the following scheme.
          1. CPU is the parameter server and manages gradient updates.
          2. Parameters are distributed evenly across all GPUs, and the first GPU
               manages gradient updates.
  
          Args:
              features: a list of tensors, one for each tower
              labels: a list of tensors, one for each tower
              mode: ModeKeys.TRAIN or EVAL
              params: Hyperparameters suitable for tuning
          Returns:
              A EstimatorSpec object.
          """ is_training = ( mode == tf . estimator . ModeKeys . TRAIN ) weight_decay = params . weight_decay momentum = params . momentum tower_features = features tower_labels = labels tower_losses = [ ] tower_gradvars = [ ] tower_preds = [ ] # channels first (NCHW) is normally optimal on GPU and channels last (NHWC) # on CPU. The exception is Intel MKL on CPU which is optimal with # channels_last. data_format = params . data_format if not data_format : if num_gpus == 0 : data_format = 'channels_last' else : data_format = 'channels_first' if num_gpus == 0 : num_devices = 1 device_type = 'cpu' else : num_devices = num_gpus device_type = 'gpu' for i in range ( num_devices ) : worker_device = '/{}:{}' . format ( device_type , i ) if variable_strategy == 'CPU' : device_setter = mlp_utils . local_device_setter ( worker_device = worker_device ) elif variable_strategy == 'GPU' : device_setter = mlp_utils . local_device_setter ( ps_device_type = 'gpu' , worker_device = worker_device , ps_strategy = tf . contrib . training . GreedyLoadBalancingStrategy ( num_gpus , tf . contrib . training . byte_size_load_fn ) ) with tf . variable_scope ( 'mlp' , reuse = bool ( i != 0 ) ) : with tf . name_scope ( 'tower_%d' % i ) as name_scope : with tf . device ( device_setter ) : loss , gradvars , preds = _tower_fn ( is_training , weight_decay , tower_features [ i ] , tower_labels [ i ] , data_format , params . batch_norm_decay , params . batch_norm_epsilon ) tower_losses . append ( loss ) tower_gradvars . append ( gradvars ) tower_preds . append ( preds ) if i == 0 : # Only trigger batch_norm moving mean and variance update from # the 1st tower. Ideally, we should grab the updates from all # towers but these stats accumulate extremely fast so we can # ignore the other stats from the other towers without # significant detriment. update_ops = tf . get_collection ( tf . GraphKeys . UPDATE_OPS , name_scope ) # Now compute global loss and gradients. gradvars = [ ] with tf . name_scope ( 'gradient_averaging' ) : all_grads = { } for grad , var in itertools . chain ( * tower_gradvars ) : if grad is not None : all_grads . setdefault ( var , [ ] ) . append ( grad ) for var , grads in six . iteritems ( all_grads ) : # Average gradients on the same device as the variables # to which they apply. with tf . device ( var . device ) : if len ( grads ) == 1 : avg_grad = grads [ 0 ] else : avg_grad = tf . multiply ( tf . add_n ( grads ) , 1. / len ( grads ) ) gradvars . append ( ( avg_grad , var ) ) # Device that runs the ops to apply global gradient updates. consolidation_device = '/gpu:0' if variable_strategy == 'GPU' else '/cpu:0' with tf . device ( consolidation_device ) : num_batches_per_epoch = mlp_data . MlpDataSet . num_examples_per_epoch ( 'train' ) // ( params . train_batch_size * num_workers ) boundaries = [ num_batches_per_epoch * x for x in np . array ( [ 20 , 50 , 80 ] , dtype = np . int64 ) ] staged_lr = [ params . learning_rate * x for x in [ 1 , 0.1 , 0.01 , 0.002 ] ] learning_rate = tf . train . piecewise_constant ( tf . train . get_global_step ( ) , boundaries , staged_lr ) loss = tf . reduce_mean ( tower_losses , name = 'loss' ) examples_sec_hook = mlp_utils . ExamplesPerSecondHook ( params . train_batch_size , every_n_steps = 10 ) tensors_to_log = { 'learning_rate' : learning_rate , 'loss' : loss } logging_hook = tf . train . LoggingTensorHook ( tensors = tensors_to_log , every_n_iter = 100 ) train_hooks = [ logging_hook , examples_sec_hook ] optimizer = tf . train . MomentumOptimizer ( learning_rate = learning_rate , momentum = momentum ) if params . sync : optimizer = tf . train . SyncReplicasOptimizer ( optimizer , replicas_to_aggregate = num_workers ) sync_replicas_hook = optimizer . make_session_run_hook ( params . is_chief ) train_hooks . append ( sync_replicas_hook ) # Create single grouped train op train_op = [ optimizer . apply_gradients ( gradvars , global_step = tf . train . get_global_step ( ) ) ] train_op . extend ( update_ops ) train_op = tf . group ( * train_op ) predictions = { 'classes' : tf . concat ( [ p [ 'classes' ] for p in tower_preds ] , axis = 0 ) , 'probabilities' : tf . concat ( [ p [ 'probabilities' ] for p in tower_preds ] , axis = 0 ) } stacked_labels = tf . concat ( labels , axis = 0 ) metrics = { 'accuracy' : tf . metrics . accuracy ( stacked_labels , predictions [ 'classes' ] ) } return tf . estimator . EstimatorSpec ( mode = mode , predictions = predictions , loss = loss , train_op = train_op , training_hooks = train_hooks , eval_metric_ops = metrics ) return _mlp_model_fn