Method apply_gradients

tensorpack/tfutils/optimizer.py:179–224

(self, grads_and_vars, global_step=None, name=None)

    @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):
                g, v = gv
                ops.append(s.assign_add(g))
            update_counter = tfv1.assign_add(counter, 1, name='update_counter')
            update_slot_op = tf.group(update_counter, *ops, name='update_slot')

            def update_grad():
                update_op = self._opt.apply_gradients(slots_and_vars)
                with tf.control_dependencies([update_op]):
                    clear_ops = [tfv1.assign(s, tf.zeros_like(s)) for s in slots]
                return tf.group(*clear_ops, name='update_grad')

            pred = tf.equal(tfv1.mod(counter, self._niter), 0)
            with tf.control_dependencies([update_slot_op]):
                if name is None:
                    name = 'cond_update_grad'
                op = tf.cond(pred, update_grad, tf.no_op)

            if global_step is not None:
                # Tensorpack maintains global_step by other means,
                # so this option is useless in tensorpack trainers.
                # But we include the implementation here for completeness
                global_step_increment = tfv1.assign_add(global_step, 1)
                op = tf.group(op, global_step_increment, name=name)
            else:
                op = tf.identity(op, name=name).op
        return op


if __name__ == '__main__':