Class ScaleGradient

tensorpack/tfutils/gradproc.py:221–264

Scale certain gradient by a multiplier.

class ScaleGradient(MapGradient):
    """
    Scale certain gradient by a multiplier.
    """

    def __init__(self, multipliers, verbose=True):
        """
        Args:
            multipliers (tuple or list): tuple of (regex, float), or list of such tuples.
            verbose (bool): whether to print logs or not

        Example:
            Use double learning rate for all the bias (as in caffe), and freeze layer0:

            .. code-block:: python

                from tensorpack.tfutils import optimizer, gradproc
                opt = optimizer.apply_grad_processors(
                    opt, [gradproc.ScaleGradient(
                        [('.*/b', 2.), ('layer0/.*', 0.)]
                    )])
        """
        if not isinstance(multipliers, list):
            multipliers = [multipliers]
        self.multipliers = multipliers
        assert verbose in [True, False], verbose
        self._verbose = verbose
        super(ScaleGradient, self).__init__(self._mapper)

    def _mapper(self, grad, var):
        varname = var.op.name
        for regex, val in self.multipliers:
            # always match against the whole name
            if not regex.endswith('$'):
                regex = regex + '$'

            if re.match(regex, varname):
                if self._verbose:
                    logger.info("Gradient of '{}' is multipled by {}".format(varname, val))
                if val != 0:    # skip zero to speed up
                    return grad * val
                else:
                    return None
        return grad