Python tensorflow.custom_gradient() Examples

def get_loss_fn(self, minimization_problem):
    """Returns the loss function.

    The resulting loss function should use `tf.custom_gradient` to override its
    gradients. First, the gradients w.r.t. the internal state should be written
    in terms of the constraints, instead of the proxy_constraints. Second, the
    gradients may be negated, depending on the formulation (for example, for the
    Lagrangian formulation, we wish to maximize over the Lagrange multipliers,
    so the associated gradients will be negated).

    Args:
      minimization_problem: `ConstrainedMinimizationProblem`, the problem to
        minimize.

    Returns:
      The loss function.
    """ 

def quantize(self, tensor):
        @tf.custom_gradient
        def op(tensor):
            def grad(dy):
                return dy
            if self.auto_threshold:
                threshold = -0.7 * tf.math.reduce_sum(tf.math.abs(tensor))/tf.dtypes.cast(tf.size(tensor),tensor.dtype)
            else:
                threshold = self.threshold
            out = tf.ones_like(tensor)*-1
            out += tf.dtypes.cast(tf.greater(tensor, -threshold),tensor.dtype)
            out += tf.dtypes.cast(tf.greater(tensor, threshold),tensor.dtype)
            out *= self.marginal
            out = tf.identity(out, name=str(self)+"_output")
            return out, grad
        return op(tensor)



###############################
### Other
############################### 

def __call__(self, *parameters, solver_args={}):
        """Solve problem (or a batch of problems) corresponding to `parameters`

        Args:
          parameters: a sequence of tf.Tensors; the n-th Tensor specifies
                      the value for the n-th CVXPY Parameter. These Tensors
                      can be batched: if a Tensor has 3 dimensions, then its
                      first dimension is interpreted as the batch size.
          solver_args: a dict of optional arguments, to send to `diffcp`. Keys
                       should be the names of keyword arguments.

        Returns:
          a list of optimal variable values, one for each CVXPY Variable
          supplied to the constructor.
        """
        if len(parameters) != len(self.params):
            raise ValueError('A tensor must be provided for each CVXPY '
                             'parameter; received %d tensors, expected %d' % (
                                 len(parameters), len(self.params)))
        compute = tf.custom_gradient(
            lambda *parameters: self._compute(parameters, solver_args))
        return compute(*parameters) 

def P_quantize(self, tensor):
        @tf.custom_gradient
        def op(tensor):
            def grad(dy):
                return dy
            #randn = tf.random.uniform(tensor.shape, minval=0, maxval=1 )
            #mask = tf.dtypes.cast(tf.less(tensor,randn), tensor.dtype)
            out= tf.math.floor(tf.math.log(tf.math.abs(tensor))/tf.math.log(tf.constant(2,dtype=tensor.dtype))) #+ mask
            out= tf.math.pow(2*tf.ones_like(tensor),out)
            out= out*tf.sign(tensor)
            out = tf.identity(out, name=str(self)+"_output")
            return out, grad
        return op(tensor)


###############################
### Sparse
############################### 

def quantize(self,tensor):
        @tf.custom_gradient
        def op(tensor):
            def grad(dy):
                return dy
            randn = tf.random.uniform(tensor.shape, minval=0, maxval=1 )
            out_up = tf.math.ceil( tensor*(1<<self.fixed_prec) ) / (1<<self.fixed_prec)
            out_down = tf.math.floor( tensor*(1<<self.fixed_prec) ) / (1<<self.fixed_prec)
            out_mask = tf.less_equal( (tensor-tf.math.floor(tensor))*(1<<self.fixed_prec) ,randn )
            out = out_down * tf.dtypes.cast(out_mask, tensor.dtype) + out_up * tf.dtypes.cast(tf.math.logical_not(out_mask), tensor.dtype)
            # handle overflow (saturate number towards maximum or minimum)
            out = tf.math.maximum( tf.math.minimum( out, self.fixed_max_signed ), self.fixed_min_signed)
            # tag output
            out = tf.identity(out, name=str(self)+"_output")
            return out, grad
        return op(tensor)


###############################
### Logarithmic
############################### 

def ste_tern(
    x: tf.Tensor,
    threshold_value: float = 0.05,
    ternary_weight_networks: bool = False,
    clip_value: float = 1.0,
) -> tf.Tensor:
    @tf.custom_gradient
    def _call(x):
        if ternary_weight_networks:
            threshold = 0.7 * tf.reduce_sum(tf.abs(x)) / tf.cast(tf.size(x), x.dtype)
        else:
            threshold = threshold_value

        def grad(dy):
            return _clipped_gradient(x, dy, clip_value)

        return tf.sign(tf.sign(x + threshold) + tf.sign(x - threshold)), grad

    return _call(x) 

def __init__(self, name, inputs, tower_setup, initial_weights, hack_gradient_magnitude=1.0):
    super().__init__()
    assert len(initial_weights) == len(inputs)
    with tf.variable_scope(name):
      initializer = tf.constant_initializer(initial_weights)
      weights = self.create_bias_variable("linear_combination_weights", len(inputs), tower_setup,
                                          initializer=initializer)
      if hack_gradient_magnitude > 1.0:
        # https://stackoverflow.com/a/43948872
        @tf.custom_gradient
        def amplify_gradient_layer(x):
          def grad(dy):
            return hack_gradient_magnitude * dy
          return tf.identity(x), grad
        weights = amplify_gradient_layer(weights)
      y = inputs[0] * weights[0]
      for n in range(1, len(inputs)):
        y += inputs[n] * weights[n]
      self.outputs.append(y)
      for n in range(len(inputs)):
        self.add_scalar_summary(weights[n], "linear_combination_weights_" + str(n)) 

def eager_tensorflow_from_torch(func):
    """
    Wraps a PyTorch function into a TensorFlow eager-mode function (ie can be executed within Tensorflow eager-mode).

    :param func: Function that takes PyTorch tensors and returns a PyTorch tensor.
    :return: Differentiable Tensorflow eager-mode function.
    """

    @tf.custom_gradient
    def compute(*inputs):
        th_inputs = [th.tensor(tf_input.numpy(), requires_grad=True) for tf_input in inputs]
        th_output = func(*th_inputs)

        def compute_grad(d_output):
            th_d_output = th.tensor(d_output.numpy(), requires_grad=False)
            th_gradients = th.autograd.grad([th_output], th_inputs, grad_outputs=[th_d_output], allow_unused=True)
            tf_gradients = [tf.convert_to_tensor(th_gradient.numpy()) for th_gradient in th_gradients]
            return tf_gradients

        return tf.convert_to_tensor(th_output.detach().numpy()), compute_grad

    return compute 

def downscale2d(x, factor=2):
    with tf.variable_scope('Downscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _downscale2d(x, factor)
            @tf.custom_gradient
            def grad(dy):
                dx = _upscale2d(dy, factor, gain=1/factor**2)
                return dx, lambda ddx: _downscale2d(ddx, factor)
            return y, grad
        return func(x)

#----------------------------------------------------------------------------
# Get/create weight tensor for a convolutional or fully-connected layer. 

def blur2d(x, f, normalize=True):
    with tf.variable_scope('Blur2D'):
        @tf.custom_gradient
        def func(x):
            y = _blur2d(x, f, normalize)

            @tf.custom_gradient
            def grad(dy):
                dx = _blur2d(dy, f, normalize, flip=True)
                return dx, lambda ddx: _blur2d(ddx, f, normalize)

            return y, grad

        return func(x) 

def apply_binomial_filter(x, f=[1,2,1], normalize=True):
    with tf.variable_scope('Blur2D'):
        @tf.custom_gradient
        def func(x):
            y = apply_binomial_filter_(x)
            @tf.custom_gradient
            def grad(dy):
                def gradgrad(ddx):
                    return apply_binomial_filter_(ddx)
                dx = apply_binomial_filter_(dy)
                return dx, gradgrad
            return y, grad
        return func(x) 

def ternarize(x, thresh=0.05):
    """
    Implemented Trained Ternary Quantization:
    https://arxiv.org/abs/1612.01064

    Code modified from the authors' at:
    https://github.com/czhu95/ternarynet/blob/master/examples/Ternary-Net/ternary.py
    """
    shape = x.get_shape()

    thre_x = tf.stop_gradient(tf.reduce_max(tf.abs(x)) * thresh)

    w_p = tf.get_variable('Wp', initializer=1.0, dtype=tf.float32)
    w_n = tf.get_variable('Wn', initializer=1.0, dtype=tf.float32)

    tf.summary.scalar(w_p.op.name + '-summary', w_p)
    tf.summary.scalar(w_n.op.name + '-summary', w_n)

    mask = tf.ones(shape)
    mask_p = tf.where(x > thre_x, tf.ones(shape) * w_p, mask)
    mask_np = tf.where(x < -thre_x, tf.ones(shape) * w_n, mask_p)
    mask_z = tf.where((x < thre_x) & (x > - thre_x), tf.zeros(shape), mask)

    @tf.custom_gradient
    def _sign_mask(x):
        return tf.sign(x) * mask_z, lambda dy: dy

    w = _sign_mask(x)

    w = w * mask_np

    tf.summary.histogram(w.name, w)
    return w 

def leaky_relu(x, alpha=0.2):
    with tf.variable_scope('LeakyReLU'):
        alpha = tf.constant(alpha, dtype=x.dtype, name='alpha')
        @tf.custom_gradient
        def func(x):
            y = tf.maximum(x, x * alpha)
            @tf.custom_gradient
            def grad(dy):
                dx = tf.where(y >= 0, dy, dy * alpha)
                return dx, lambda ddx: tf.where(y >= 0, ddx, ddx * alpha)
            return y, grad
        return func(x)

#----------------------------------------------------------------------------
# Pixelwise feature vector normalization. 

def TensorflowDifferentiableRender2D(ij, colors, scene):
    """Tensorflow implementation of the 2D rendering function."""

    @tf.custom_gradient
    def forward(
        ij, colors
    ):  # using inner function as we don't differentate w.r.t scene
        nb_color_chanels = colors.shape[1]
        image = np.empty((scene.height, scene.width, nb_color_chanels))
        z_buffer = np.empty((scene.height, scene.width))
        scene.ij = np.array(ij)  # should automatically detached according to
        # https://pytorch.org/docs/master/notes/extending.html
        scene.colors = np.array(colors)
        scene.depths = np.array(scene.depths)
        differentiable_renderer_cython.renderScene(scene, 1, image, z_buffer)

        def backward(image_b):
            scene.uv_b = np.zeros(scene.uv.shape)
            scene.ij_b = np.zeros(scene.ij.shape)
            scene.shade_b = np.zeros(scene.shade.shape)
            scene.colors_b = np.zeros(scene.colors.shape)
            scene.texture_b = np.zeros(scene.texture.shape)
            image_copy = (
                image.copy()
            )  # making a copy to avoid removing antialiasing on the image returned by
            # the forward pass (the c++ backpropagation undo antialiasing), could be
            # optional if we don't care about getting aliased images
            differentiable_renderer_cython.renderSceneB(
                scene, 1, image_copy, z_buffer, image_b.numpy()
            )
            return tf.constant(scene.ij_b), tf.constant(scene.colors_b)

        return tf.convert_to_tensor(image), backward

    return forward(ij, colors) 

def blur2d(x, f=[1,2,1], normalize=True):
    with tf.variable_scope('Blur2D'):
        @tf.custom_gradient
        def func(x):
            y = _blur2d(x, f, normalize)
            @tf.custom_gradient
            def grad(dy):
                dx = _blur2d(dy, f, normalize, flip=True)
                return dx, lambda ddx: _blur2d(ddx, f, normalize)
            return y, grad
        return func(x) 

def upscale2d(x, factor=2):
    with tf.variable_scope('Upscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _upscale2d(x, factor)
            @tf.custom_gradient
            def grad(dy):
                dx = _downscale2d(dy, factor, gain=factor**2)
                return dx, lambda ddx: _upscale2d(ddx, factor)
            return y, grad
        return func(x) 

def blur2d(x, f=[1,2,1], normalize=True):
    with tf.variable_scope('Blur2D'):
        @tf.custom_gradient
        def func(x):
            y = _blur2d(x, f, normalize)
            @tf.custom_gradient
            def grad(dy):
                dx = _blur2d(dy, f, normalize, flip=True)
                return dx, lambda ddx: _blur2d(ddx, f, normalize)
            return y, grad
        return func(x) 

def leaky_relu(x, alpha=0.2):
    with tf.variable_scope('LeakyReLU'):
        alpha = tf.constant(alpha, dtype=x.dtype, name='alpha')
        @tf.custom_gradient
        def func(x):
            y = tf.maximum(x, x * alpha)
            @tf.custom_gradient
            def grad(dy):
                dx = tf.where(y >= 0, dy, dy * alpha)
                return dx, lambda ddx: tf.where(y >= 0, ddx, ddx * alpha)
            return y, grad
        return func(x)

#----------------------------------------------------------------------------
# Pixelwise feature vector normalization. 

def _upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
    """Fast CUDA implementation of `upfirdn_2d()` using custom ops."""

    x = tf.convert_to_tensor(x)
    k = np.asarray(k, dtype=np.float32)
    majorDim, inH, inW, minorDim = x.shape.as_list()
    kernelH, kernelW = k.shape
    assert inW >= 1 and inH >= 1
    assert kernelW >= 1 and kernelH >= 1
    assert isinstance(upx, int) and isinstance(upy, int)
    assert isinstance(downx, int) and isinstance(downy, int)
    assert isinstance(padx0, int) and isinstance(padx1, int)
    assert isinstance(pady0, int) and isinstance(pady1, int)

    outW = (inW * upx + padx0 + padx1 - kernelW) // downx + 1
    outH = (inH * upy + pady0 + pady1 - kernelH) // downy + 1
    assert outW >= 1 and outH >= 1

    kc = tf.constant(k, dtype=x.dtype)
    gkc = tf.constant(k[::-1, ::-1], dtype=x.dtype)
    gpadx0 = kernelW - padx0 - 1
    gpady0 = kernelH - pady0 - 1
    gpadx1 = inW * upx - outW * downx + padx0 - upx + 1
    gpady1 = inH * upy - outH * downy + pady0 - upy + 1

    @tf.custom_gradient
    def func(x):
        y = _get_plugin().up_fir_dn2d(x=x, k=kc, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
        y.set_shape([majorDim, outH, outW, minorDim])
        @tf.custom_gradient
        def grad(dy):
            dx = _get_plugin().up_fir_dn2d(x=dy, k=gkc, upx=downx, upy=downy, downx=upx, downy=upy, padx0=gpadx0, padx1=gpadx1, pady0=gpady0, pady1=gpady1)
            dx.set_shape([majorDim, inH, inW, minorDim])
            return dx, func
        return y, grad
    return func(x)

#---------------------------------------------------------------------------- 

def upscale2d(x, factor=2):
    with tf.variable_scope('Upscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _upscale2d(x, factor)
            @tf.custom_gradient
            def grad(dy):
                dx = _downscale2d(dy, factor, gain=factor**2)
                return dx, lambda ddx: _upscale2d(ddx, factor)
            return y, grad
        return func(x) 

def downscale2d(x, factor=2):
    with tf.variable_scope('Downscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _downscale2d(x, factor)
            @tf.custom_gradient
            def grad(dy):
                dx = _upscale2d(dy, factor, gain=1/factor**2)
                return dx, lambda ddx: _downscale2d(ddx, factor)
            return y, grad
        return func(x)

#----------------------------------------------------------------------------
# Get/create weight tensor for a convolutional or fully-connected layer. 

def downscale2d(x, factor=2):
    with tf.variable_scope('Downscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _downscale2d(x, factor)
            @tf.custom_gradient
            def grad(dy):
                dx = _upscale2d(dy, factor, gain=1/factor**2)
                return dx, lambda ddx: _downscale2d(ddx, factor)
            return y, grad
        return func(x)

#----------------------------------------------------------------------------
# Get/create weight tensor for a convolutional or fully-connected layer. 

def leaky_relu(x, alpha=0.2):
    with tf.variable_scope('LeakyReLU'):
        alpha = tf.constant(alpha, dtype=x.dtype, name='alpha')
        @tf.custom_gradient
        def func(x):
            y = tf.maximum(x, x * alpha)
            @tf.custom_gradient
            def grad(dy):
                dx = tf.where(y >= 0, dy, dy * alpha)
                return dx, lambda ddx: tf.where(y >= 0, ddx, ddx * alpha)
            return y, grad
        return func(x)

#----------------------------------------------------------------------------
# Pixelwise feature vector normalization. 

def upscale2d(x, factor=2):
    with tf.variable_scope('Upscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _upscale2d(x, factor)
            @tf.custom_gradient
            def grad(dy):
                dx = _downscale2d(dy, factor, gain=factor**2)
                return dx, lambda ddx: _upscale2d(ddx, factor)
            return y, grad
        return func(x) 

def quantize(self, tensor):
        @tf.custom_gradient
        def op(tensor):
            def grad(dy):
                return dy
            out=tf.dtypes.cast(tf.dtypes.cast(tensor,tf.float16),tensor.dtype)
            out = tf.identity(out, name=str(self)+"_output")
            return out, grad
        return op(tensor)

###############################
### Binary
############################### 

def upscale2d(x, factor=2):
    with tf.variable_scope('Upscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _upscale2d(x, factor)

            @tf.custom_gradient
            def grad(dy):
                dx = _downscale2d(dy, factor, gain=factor ** 2)
                return dx, lambda ddx: _upscale2d(ddx, factor)

            return y, grad

        return func(x) 

def downscale2d(x, factor=2):
    with tf.variable_scope('Downscale2D'):
        @tf.custom_gradient
        def func(x):
            y = _downscale2d(x, factor)

            @tf.custom_gradient
            def grad(dy):
                dx = _upscale2d(dy, factor, gain=1 / factor ** 2)
                return dx, lambda ddx: _downscale2d(ddx, factor)

            return y, grad

        return func(x) 

def call(self, input_):
        @tf.custom_gradient
        def _call(input_):
            def reversed_gradient(output_grads):
                return self.weight * tf.negative(output_grads)

            return input_, reversed_gradient
        
        return _call(input_)


# ## The model function
# The network consists of 3 sub-networks:
#
# * Feature extractor: extracts internal representation for both the source and target distributions.
#
# * Label predictor: predicts label from the extracted features.
#
# * Domain classifier: classifies the origin (`source` or `target`) of the extracted features.
#
#
# Both the label predictor and the domain classifier will try to minimize 
# classification loss, but the gradients backpropagated from the domain
# classifier to the feature extractor have their signs reversed.
#
#
# This model function also shows how to use `host_call` to output summaries.
# 

def edge_softmax(graph, logits, eids=ALL):
    """Closure for tf.custom_gradient"""

    @tf.custom_gradient
    def _lambda(logits):
        return edge_softmax_real(graph, logits, eids=eids)

    return _lambda(logits) 

def binary_reduce(reducer, binary_op, graph, lhs, rhs, lhs_data, rhs_data,
                  out_size, lhs_map=(None, None), rhs_map=(None, None), out_map=(None, None)):

    @tf.custom_gradient
    def _lambda(lhs_data, rhs_data):
        return binary_reduce_real(reducer, binary_op, graph, lhs, rhs, lhs_data, rhs_data,
                                  out_size, lhs_map, rhs_map, out_map)
    return _lambda(lhs_data, rhs_data)