tensorflow 中 ASGD with Delay Compensation 优化器代码实现

"""DelayCompensatedGradientDescentOptimizer for TensorFlow."""from __future__ import absolute_importfrom __future__ import divisionfrom __future__ import print_function
from tensorflow.python.framework import opsfrom tensorflow.python.ops import array_opsfrom tensorflow.python.ops import control_flow_opsfrom tensorflow.python.ops import math_opsfrom tensorflow.python.ops import state_opsfrom tensorflow.python.training import optimizerfrom tensorflow.python.training import training_ops
GATE_NONE = 0GATE_OP = 1GATE_GRAPH = 2

class DelayCompensatedGradientDescentOptimizer(optimizer.Optimizer):    """Optimizer that implements the DelayCompensatedGradientDescent algorithm.    See [](https://arxiv.org/abs/1609.08326)    ([](https://arxiv.org/pdf/1609.08326.pdf)).    """
    def __init__(self, learning_rate, variance_parameter=2.0, num_workers=1,                 use_locking=False, name="DelayCompensatedGradientDescentOptimizer"):
        """Construct a gradient descent optimizer with delay compensation.        It is cricial to note the `num_workers` in constructor and `worker_index` in        `minimize()` and `apply_gradients()`.        Contrast to AdaMaxamOptimizer, the sparse implementation of this algorithm        (used when the gradient is an IndexedSlices object, typically because of        `tf.gather` or an embedding lookup in the forward pass) only updates        variable slices and corresponding `shadow_t` term when that part of        the variable was used in the forward pass. This means that the sparse        behavior is contrast to the dense behavior (similar to some momentum        implementations which ignore momentum unless a variable slice was actually        used).        Args:          learning_rate: A Tensor or a floating point value. The learning rate.          variance_parameter: A Tensor or a floating point value.            The variance control parameter.          num_workers: A int value. The number of workers.          use_locking: If True use locks for update operations.          name: Optional name for the operations created when applying gradients.            Defaults to "DelayCompensatedGradientDescentOptimizer".        """        num_workers = self._call_if_callable(num_workers)        if num_workers <= 0:            raise ValueError("num_workers must be positive: %s" % num_workers)        super(DelayCompensatedGradientDescentOptimizer, self).__init__(use_locking, name)        self._lr = learning_rate        self._lambda = variance_parameter        self._num_workers = num_workers        self._learning_rate_tensor = None        self._lambda_tensor = None        self._use_locking = use_locking
    def _create_slots(self, var_list):        for index in range(self._num_workers):            for v in var_list:                self._zeros_slot(v, "shadow_{0}".format(index), self._name)
    def _prepare(self):        lr = self._call_if_callable(self._lr)        lambda_ = self._call_if_callable(self._lambda)
        self._learning_rate_tensor = ops.convert_to_tensor(lr, name="learning_rate")        self._lambda_tensor = ops.convert_to_tensor(lambda_, name="lambda")
    def _apply_dense(self, grad, var):
        shadow = self.get_slot(var, "shadow_{0}".format(self.worker_index))        return training_ops.apply_delay_compensated_gradient_descent(            var,            math_ops.cast(self._learning_rate_tensor, grad.dtype.base_dtype),            grad,            math_ops.cast(self._lambda_tensor, grad.dtype.base_dtype),            shadow,            use_locking=self._use_locking).op
    def _resource_apply_dense(self, grad, var):
        shadow = self.get_slot(var, "shadow_{0}".format(self.worker_index))        return training_ops.resource_apply_delay_compensated_gradient_descent(            var.handle,            math_ops.cast(self._learning_rate_tensor, grad.dtype.base_dtype),            grad,            math_ops.cast(self._lambda_tensor, grad.dtype.base_dtype),            shadow.handle,            use_locking=self._use_locking)
    def _apply_sparse_shared(self, grad, var, indices):
        shadow = self.get_slot(var, "shadow_{0}".format(self.worker_index))        # if shadow is None:        #   raise ValueError("None shadow with index = " + str(self.worker_index) + " and var = " + str(var))        lambda_ = math_ops.cast(self._lambda_tensor, var.dtype.base_dtype)        lr = math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype)
        var_slice = array_ops.gather(var, indices)        shadow_slice = array_ops.gather(shadow, indices)
        var_scaled_g_values = lr * (grad + lambda_ * grad * grad * (var_slice - shadow_slice))
        var_t = state_ops.scatter_add(var, indices, -var_scaled_g_values, use_locking=self._use_locking)
        with ops.control_dependencies([var_t]):            shadow_t = state_ops.assign(shadow, var_t)
        return control_flow_ops.group(*[var_t, shadow_t])
    def _apply_sparse(self, grad, var):        return self._apply_sparse_shared(            grad.values, var, grad.indices)
    def _resource_apply_sparse(self, grad, var, indices):        return self._apply_sparse_shared(            grad, var, indices)
    def minimize(self, loss, global_step=None, var_list=None,                 gate_gradients=GATE_OP, aggregation_method=None,                 colocate_gradients_with_ops=False, name=None,                 grad_loss=None, worker_index=0):        self.worker_index = worker_index        return super(DelayCompensatedGradientDescentOptimizer, self).minimize(loss=loss, global_step=global_step,                                                                              var_list=var_list,                                                                              gate_gradients=gate_gradients,                                                                              aggregation_method=aggregation_method,                                                                              colocate_gradients_with_ops=colocate_gradients_with_ops,                                                                              name=name,                                                                              grad_loss=grad_loss)
    def apply_gradients(self, grads_and_vars, global_step=None, name=None, worker_index=0):        self.worker_index = worker_index        return super(DelayCompensatedGradientDescentOptimizer, self).apply_gradients(grads_and_vars=grads_and_vars,

                                                                                 global_step=global_step, name=name)

import numpy as npimport tensorflow as tffrom tensorflow.python.framework import constant_opfrom tensorflow.python.framework import opsfrom tensorflow.python.ops import variablesfrom tensorflow.python.training import adam

if __name__ == '__main__':    value_a = np.ones(shape=[3, 10])    indices_a = np.array([0, 3, 8])    dense_shape_a = [10, 10]    grad_slices_a = ops.IndexedSlices(constant_op.constant(value_a), constant_op.constant(indices_a),                                      constant_op.constant(dense_shape_a))
    var_np = np.ones(shape=[10, 10])
    var0 = variables.RefVariable(var_np)    opt = adam.AdamOptimizer()    update = opt.apply_gradients(zip([grad_slices_a], [var0]))    # variables.global_variables_initializer().run()    sess = tf.Session()    sess.run(tf.global_variables_initializer())    print("initial variable is:", sess.run(var0))    sess.run(update)    print("update 1 time variable is:", sess.run(var0))

输出：initial variable is: [[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.] [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]]update 1 time variable is: [[0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999] [1.    1.    1.    1.    1.    1.    1.    1.    1.    1.   ] [1.    1.    1.    1.    1.    1.    1.    1.    1.    1.   ] [0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999] [1.    1.    1.    1.    1.    1.    1.    1.    1.    1.   ] [1.    1.    1.    1.    1.    1.    1.    1.    1.    1.   ] [1.    1.    1.    1.    1.    1.    1.    1.    1.    1.   ] [1.    1.    1.    1.    1.    1.    1.    1.    1.    1.   ] [0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999 0.999] [1.    1.    1.    1.    1.    1.    1.    1.    1.    1.   ]]

template <typename T>struct ApplyDelayCompensatedGradientDescent<CPUDevice, T> {  void operator()(const CPUDevice& d, typename TTypes<T>::Flat var,                   typename TTypes<T>::ConstScalar lr,                   typename TTypes<T>::ConstFlat grad,                   typename TTypes<T>::ConstScalar variance,                   typename TTypes<T>::Flat shadow) {    var.device(d) -= lr() * (grad + variance() * grad * grad * (var - shadow));    shadow.device(d) = var;  }};

from __future__ import print_function, absolute_import, division
import tensorflow as tf
tf.app.flags.DEFINE_string("ps_hosts", "localhost:2222", "ps hosts")tf.app.flags.DEFINE_string("worker_hosts", "localhost:2223,localhost:2224", "worker hosts")tf.app.flags.DEFINE_string("job_name", "worker", "'ps' or'worker'")tf.app.flags.DEFINE_integer("task_index", 0, "Index of task within the job")tf.app.flags.DEFINE_integer("num_workers", 2, "Number of workers")tf.app.flags.DEFINE_boolean("is_sync", False, "using synchronous training or not")
FLAGS = tf.app.flags.FLAGS

def model(images):    """Define a simple mnist classifier"""    net = tf.layers.dense(images, 500, activation=tf.nn.relu)    net = tf.layers.dense(net, 500, activation=tf.nn.relu)    net = tf.layers.dense(net, 10, activation=None)    return net

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()x_train = x_train.reshape(-1, 784).astype('float32')x_test = x_test.reshape(-1, 784).astype('float32')x_train /= 255x_test /= 255

def get_batch(image, label, batch_size=32, training=True):    df = tf.data.Dataset.from_tensor_slices((image, label))    if training:        df = df.repeat(10).shuffle(buffer_size=1000)    df = df.batch(batch_size).prefetch(batch_size)    iterator = df.make_one_shot_iterator()    batch_x, batch_y = iterator.get_next()    return batch_x, batch_y

def main(_):    ps_hosts = FLAGS.ps_hosts.split(",")    worker_hosts = FLAGS.worker_hosts.split(",")
    # create the cluster configured by `ps_hosts' and 'worker_hosts'    cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
    # create a server for local task    server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)
    train_batch_x, train_batch_y = get_batch(x_train, y_train)    test_batch_x, test_batch_y = get_batch(x_test, y_test, training=False)
    if FLAGS.job_name == "ps":        server.join()  # ps hosts only join    elif FLAGS.job_name == "worker":        # workers perform the operation        # ps_strategy = tf.contrib.training.GreedyLoadBalancingStrategy(FLAGS.num_ps)
        # Note: tf.train.replica_device_setter automatically place the paramters (Variables)        # on the ps hosts (default placement strategy:  round-robin over all ps hosts, and also        # place multi copies of operations to each worker host        with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % FLAGS.task_index,                                                      cluster=cluster)):
            logits = model(train_batch_x)            loss = tf.reduce_mean(                tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=tf.one_hot(train_batch_y, 10)))
            # The StopAtStepHook handles stopping after running given steps.            hooks = [tf.train.StopAtStepHook(last_step=10000)]
            global_step = tf.train.get_or_create_global_step()            #optimizer = tf.train.AdamOptimizer(learning_rate=1e-04)            optimizer = tf.contrib.opt.DelayCompensatedGradientDescentOptimizer(learning_rate=0.001)
            if FLAGS.is_sync:                # asynchronous training                # use tf.train.SyncReplicasOptimizer wrap optimizer                # ref: https://www.tensorflow.org/api_docs/python/tf/train/SyncReplicasOptimizer                optimizer = tf.train.SyncReplicasOptimizer(optimizer, replicas_to_aggregate=FLAGS.num_workers,                                                           total_num_replicas=FLAGS.num_workers)                # create the hook which handles initialization and queues                hooks.append(optimizer.make_session_run_hook((FLAGS.task_index == 0)))
            train_op = optimizer.minimize(loss, global_step=global_step)
            # The MonitoredTrainingSession takes care of session initialization,            # restoring from a checkpoint, saving to a checkpoint, and closing when done            # or an error occurs.            with tf.train.MonitoredTrainingSession(master=server.target,                                                   is_chief=(FLAGS.task_index == 0),                                                   checkpoint_dir="./checkpoint_dir",                                                   hooks=hooks) as mon_sess:                while not mon_sess.should_stop():                    # mon_sess.run handles AbortedError in case of preempted PS.                    _, ls, step = mon_sess.run([train_op, loss, global_step])                    if step % 100 == 0:                        print("Train step %d, loss: %f" % (step, ls))

if __name__ == "__main__":    tf.app.run()

python dc_asgd_exp.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224 --job_name=ps --task_index=0python dc_asgd_exp.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224 --job_name=worker --task_index=0