Class MapGradient

tensorpack/tfutils/gradproc.py:105–149

Apply a function on all gradient if the name matches regex. Keep the other gradients unchanged. It can be used for gradient clipping, etc.

class MapGradient(GradientProcessor):
    """
    Apply a function on all gradient if the name matches regex.
    Keep the other gradients unchanged.

    It can be used for gradient clipping, etc.
    """

    def __init__(self, func, regex='.*'):
        """
        Args:
            func: a user-supplied function which takes one or two arguments.
                The argument(s) can be either a `grad` tensor, or `grad` and `var`.
                The function should return the new gradient to be used.
                If it return None, the gradient is discarded (hence no update to the variable will happen).
            regex (str): used to match variables. Defaults to match all variables.
        """
        args = inspect.getfullargspec(func).args
        arg_num = len(args) - inspect.ismethod(func)
        assert arg_num in [1, 2], \
            "The function must take 1 or 2 arguments!  ({})".format(args)
        if arg_num == 1:
            self.func = lambda grad, var: func(grad)
        else:
            self.func = func

        if not regex.endswith('$'):
            regex = regex + '$'
        self.regex = regex
        super(MapGradient, self).__init__()

    def _process(self, grads):
        ret = []
        matched = False
        for grad, var in grads:
            if re.match(self.regex, var.op.name):
                matched = True
                grad = self.func(grad, var)
                if grad is not None:
                    ret.append((grad, var))
            else:
                ret.append((grad, var))
        if not matched:
            logger.warn("[MapGradient] No match was found for regex {}.".format(self.regex))
        return ret


# TODO has dependency problems: sess.run may not depend on grad