Class AccumGradOptimizer

tensorpack/tfutils/optimizer.py:141–224

An optimizer which accumulates gradients across :math:`k` :meth:`minimize` executions, and apply them together in every :math:`k` th :meth:`minimize` execution. This is roughly the same as using a :math:`k` times larger batch size plus a :math:`k` times larger learning rate, but use

class AccumGradOptimizer(ProxyOptimizer):
    """
    An optimizer which accumulates gradients across :math:`k` :meth:`minimize` executions,
    and apply them together in every :math:`k` th :meth:`minimize` execution.
    This is roughly the same as using a :math:`k` times larger batch size plus a
    :math:`k` times larger learning rate, but uses much less memory.

    This optimizer can be used in any TensorFlow code (with or without tensorpack).

    Example:

    .. code-block:: python

        from tensorpack.tfutils.optimizer import AccumGradOptimizer
        myopt = tf.train.GradientDescentOptimizer(0.01)
        myopt = AccumGradOptimizer(myopt, niter=5)
        train_op = myopt.minimize(loss)

    """

    def __init__(self, opt, niter):
        """
        Args:
            opt (tf.train.Optimizer): the underlying sub-optimizer.
            niter (int): number of iterations to accumulate gradients.
        """
        super(AccumGradOptimizer, self).__init__(opt, 'AccumGrad')
        self._niter = int(niter)

    def _create_accum_slots(self, var_list):
        slots = []
        for v in var_list:
            # TODO an option to not colocate the accumulators with variables (to save more memory)
            s = self._zeros_slot(v, "accum", self._name)
            slots.append(s)
        return slots

    @HIDE_DOC
    def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        grads_and_vars = FilterNoneGrad().process(grads_and_vars)
        vs = []
        for g, v in grads_and_vars:
            assert isinstance(g, (tf.Tensor, tf.IndexedSlices)) and isinstance(v, tf.Variable), \
                "AccumGradOptimizer does not work for the gradient of {}! " \
                "Types of v and g are {} and {}".format(v.op.name, type(v), type(g))
            vs.append(v)

        with tf.control_dependencies(None):
            slots = self._create_accum_slots(vs)
            slots_and_vars = [(s, gv[1]) for s, gv in zip(slots, grads_and_vars)]

            with tfv1.variable_scope(self._name), tf.device('/cpu:0'):
                counter = tf.Variable(
                    0, name="counter", trainable=False, dtype=tf.int32)

        with tf.name_scope('AccumGradOptimizer'):
            ops = []
            for s, gv in zip(slots, grads_and_vars):