tf.contrib.estimator.replicate_model_fn

Replicate Estimator.model_fn over GPUs. (deprecated)

tf.contrib.estimator.replicate_model_fn(
    model_fn, loss_reduction=losses.Reduction.SUM_BY_NONZERO_WEIGHTS, devices=None
)

The given model_fn specifies a single forward pass of a model. To replicate such a model over GPUs, each GPU gets its own instance of the forward pass (a.k.a. a tower). The input features and labels get sharded into the chunks that correspond to the number of GPUs. Each tower computes a loss based on its input. For each such loss, gradients are computed. After that, the available losses are aggregated to form aggregated loss. Available gradients are summed. Then, they update weights using the specified optimizer.

If devices are None, then all available GPUs are going to be used for replication. If no GPUs are available, then the model is going to be placed on the CPU.

Two modes of local replication over available GPUs are supported:

1) If exactly 1 GPU is detected, then variables and operations are placed onto the GPU. 2) If more than 1 GPU is detected, then variables are going to be placed on the CPU. Replicas of operations are placed on each individual GPU.

Here is an example of how one might use their model_fn to run over GPUs:

...
def model_fn(...):  # See `model_fn` in `Estimator`.
  loss = ...
  optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
  optimizer = tf.contrib.estimator.TowerOptimizer(optimizer)
  if mode == tf.estimator.ModeKeys.TRAIN:
    #  See the section below on <a href="/api_docs/python/tf/estimator/EstimatorSpec.md#train_op"><code>EstimatorSpec.train_op</code></a>.
    return EstimatorSpec(mode=mode, loss=loss,
                         train_op=optimizer.minimize(loss))

  #  No change for <a href="/api_docs/python/tf/estimator/ModeKeys.md#EVAL"><code>ModeKeys.EVAL</code></a> or <a href="/api_docs/python/tf/estimator/ModeKeys.md#PREDICT"><code>ModeKeys.PREDICT</code></a>.
  return EstimatorSpec(...)
...
classifier = tf.estimator.Estimator(
  model_fn=tf.contrib.estimator.replicate_model_fn(model_fn))

Please see DNNClassifierIntegrationTest for an example with a canned Estimator.

On EstimatorSpec.train_op: model_fn returns EstimatorSpec.train_op for tf.estimator.GraphKeys.TRAIN. It is typically derived using an optimizer. Towers are expected to populate it in the same way. Gradients from all towers are reduced and applied in the last tower. To achieve that in the case of multiple towers, TowerOptimizer needs to be used. See TowerOptimizer.

On sharding input features and labels: Input features and labels are split for consumption by each tower. They are split across the dimension 0. Features and labels need to be batch major.

On reduction algorithms: Certain algorithms were chosen for aggregating results of computations on multiple towers:

• Losses from all towers are reduced according to loss_reduction.
• Gradients from all towers are reduced according to loss_reduction for each trainable variable.
• eval_metrics_ops are reduced per metric using reduce_mean.
• EstimatorSpec.predictions and EstimatorSpec.export_outputs are reduced using concatenation.
• For all other fields of EstimatorSpec the values of the first tower are taken.

On distribution of variables: Variables are not duplicated between towers. Instead, they are placed on a single device as defined above and shared across towers.

On overhead:

If only one device is specified, then aggregation of loss and gradients doesn't happen. Replication consists of placing model_fn onto the specified device.

On current limitations:

• predictions are not supported for ModeKeys.EVAL. They are required for tf.contrib.estimator.add_metrics.

Args
model_fn	model_fn as defined in Estimator. See the section above about the train_op argument of EstimatorSpec.
loss_reduction	controls whether losses are summed or averaged.
devices	Optional list of devices to replicate the model across. This argument can be used to replicate only on the subset of available GPUs. If None, then all available GPUs are going to be used for replication. If no GPUs are available, then the model is going to be placed on the CPU.

Raises
ValueError	if there is no loss_reduction or if TowerOptimizer is mis-used.

Returns
A replicated version of the supplied model_fn. Returned function that conforms to the requirements of Estimator's model_fn and can be used instead of the supplied model_fn.