Python tensorflow.custom_gradient() Examples

The following are 30 code examples of tensorflow.custom_gradient(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module tensorflow , or try the search function .
Example #1
Source File: constrained_optimizer.py    From tensorflow_constrained_optimization with Apache License 2.0 6 votes vote down vote up
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.
    """ 
Example #2
Source File: Quantizers.py    From TensorQuant with Apache License 2.0 6 votes vote down vote up
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
############################### 
Example #3
Source File: cvxpylayer.py    From cvxpylayers with Apache License 2.0 6 votes vote down vote up
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) 
Example #4
Source File: Quantizers.py    From TensorQuant with Apache License 2.0 6 votes vote down vote up
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
############################### 
Example #5
Source File: Quantizers.py    From TensorQuant with Apache License 2.0 6 votes vote down vote up
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
############################### 
Example #6
Source File: quantizers.py    From larq with Apache License 2.0 6 votes vote down vote up
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) 
Example #7
Source File: UtilLayers.py    From TrackR-CNN with MIT License 6 votes vote down vote up
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)) 
Example #8
Source File: __init__.py    From tfpyth with MIT License 6 votes vote down vote up
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 
Example #9
Source File: networks_stylegan.py    From higan with MIT License 5 votes vote down vote up
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. 
Example #10
Source File: ops.py    From StyleGAN-Tensorflow with MIT License 5 votes vote down vote up
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) 
Example #11
Source File: ops.py    From stylegan_reimplementation with Apache License 2.0 5 votes vote down vote up
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) 
Example #12
Source File: dorefa.py    From tensorpack with Apache License 2.0 5 votes vote down vote up
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 
Example #13
Source File: networks_stylegan.py    From interfacegan with MIT License 5 votes vote down vote up
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. 
Example #14
Source File: differentiable_renderer_tensorflow.py    From DEODR with BSD 2-Clause "Simplified" License 5 votes vote down vote up
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) 
Example #15
Source File: networks_stylegan.py    From higan with MIT License 5 votes vote down vote up
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) 
Example #16
Source File: networks_stylegan.py    From higan with MIT License 5 votes vote down vote up
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) 
Example #17
Source File: networks_stylegan.py    From interfacegan with MIT License 5 votes vote down vote up
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) 
Example #18
Source File: networks_stylegan.py    From higan with MIT License 5 votes vote down vote up
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. 
Example #19
Source File: upfirdn_2d.py    From higan with MIT License 5 votes vote down vote up
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)

#---------------------------------------------------------------------------- 
Example #20
Source File: networks_stylegan.py    From higan with MIT License 5 votes vote down vote up
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) 
Example #21
Source File: networks_stylegan.py    From higan with MIT License 5 votes vote down vote up
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. 
Example #22
Source File: networks_stylegan.py    From interfacegan with MIT License 5 votes vote down vote up
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. 
Example #23
Source File: networks_stylegan.py    From higan with MIT License 5 votes vote down vote up
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. 
Example #24
Source File: networks_stylegan.py    From interfacegan with MIT License 5 votes vote down vote up
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) 
Example #25
Source File: Quantizers.py    From TensorQuant with Apache License 2.0 5 votes vote down vote up
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
############################### 
Example #26
Source File: ops.py    From StyleGAN-Tensorflow with MIT License 5 votes vote down vote up
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) 
Example #27
Source File: ops.py    From StyleGAN-Tensorflow with MIT License 5 votes vote down vote up
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) 
Example #28
Source File: trainer.py    From cloudml-samples with Apache License 2.0 5 votes vote down vote up
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.
# 
Example #29
Source File: softmax.py    From dgl with Apache License 2.0 5 votes vote down vote up
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) 
Example #30
Source File: tensor.py    From dgl with Apache License 2.0 5 votes vote down vote up
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)