Exluding torch.clamp() from backpropagation (as tf.stop_gradient in tensorflow)

username: always
cooked: <p>Hi,</p>
<p>when using torch.clamp(), the derivative w.r.t. to its input is zero if the input is outside [min, max]. This results in all gradients for previous operations in the graph to become zero due to the chain rule:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/b/be74cbc30b416ac044d3bf705b37dae13ecc5d6a.png" alt="20" data-base62-sha1="raQX7q2eZMxar5AvMNPpUz0pQhA" width="627" height="337"></p>
<p>In tensorflow, one can use tf.stop_gradient (<a href="https://www.tensorflow.org/api_docs/python/tf/stop_gradient" class="inline-onebox" rel="noopener nofollow ugc">tf.stop_gradient &nbsp;|&nbsp; TensorFlow v2.15.0.post1</a>) to prevent this behavior. Is there something similar for PyTorch?</p>

username: always
cooked: <p>Below is a minimal working example. A neural net with one weight is supposed to push its input towards 5. To reduce the range of possible values, the output is clipped to [4, 6].</p>
<p>Ideally, the network should still find the optimal value despite the “prior”. But it gets stuck due to zero gradient.</p>
<p>This would not happen, if the derivative of the clamp-function could be excluded from backpropagation.</p>
<pre><code class="lang-python">import torch
import numpy as np
import torch.nn as nn

x = torch.from_numpy(np.array(([1])).astype(np.float32))  # one scalar as input
layer = nn.Linear(1, 1, bias=False)  # neural net with one weight
optimizer = torch.optim.Adam(params=layer.parameters(), lr=1e-3)

for i in range(101):
    w = list(layer.parameters())[0]  # weight before backprop
    y = layer(x)  # y = w * x
    f_y = torch.clamp(y, min=4, max=6)  # f(y) = clip(y)
    loss = torch.abs(f_y - 5)   # absolute error, zero if f(y) = 5

    optimizer.zero_grad()
    loss.backward()
    grad = w.grad

    if (i % 100 == 0) or (i == 0):
        print('iteration {}'.format(i))
        print('w: {:.2f}'.format(w.detach().numpy()[0][0]))
        print('y: {:.2f}'.format(y.detach().numpy()[0]))
        print('f_y: {:.2f}'.format(f_y.detach().numpy()[0]))
        print('loss: {:.2f}'.format(loss.detach().numpy()[0]))
        print('grad: {:.2f}\n'.format(grad.detach().numpy()[0][0]))

    optimizer.step()
</code></pre>
<pre><code class="lang-auto">iteration 0
w: 0.96
y: 0.96
f_y: 4.00
loss: 1.00
grad: 0.00

iteration 100
w: 0.96
y: 0.96
f_y: 4.00
loss: 1.00
grad: 0.00
</code></pre>

username: always
cooked: <p>It turns out that the problem can be solved by creating a custom Clamp class with custom backward-method. The only remaining issue is that I do not know how to pass the min/max-values as an argument.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import numpy as np
import torch.nn as nn


class Clamp(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        return input.clamp(min=4, max=6)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output.clone()


clamp_class = Clamp()

x = torch.from_numpy(np.array(([1])).astype(np.float32))  # one scalar as input
layer = nn.Linear(1, 1, bias=False)  # neural net with one weight
optimizer = torch.optim.Adam(params=layer.parameters(), lr=1e-3)

for i in range(10001):
    w = list(layer.parameters())[0]  # weight before backprop
    y = layer(x)  # y = w * x
    clamp = clamp_class.apply
    f_y = clamp(y)  # f(y) = clip(y)
    loss = torch.abs(f_y - 5)  # absolute error, zero if f(y) = 2

    optimizer.zero_grad()
    loss.backward()
    grad = w.grad

    if (i % 100 == 0) or (i == 0):
        print('iteration {}'.format(i))
        print('w: {:.2f}'.format(w.detach().numpy()[0][0]))
        print('y: {:.2f}'.format(y.detach().numpy()[0]))
        print('f_y: {:.2f}'.format(f_y.detach().numpy()[0]))
        print('loss: {:.2f}'.format(loss.detach().numpy()[0]))
        print('grad: {:.2f}\n'.format(grad.detach().numpy()[0][0]))

    optimizer.step()
</code></pre>
<p>The plot below shows output over iterations. The output finally reaches the target value! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/c/cc58633c966c2d2cd4419de876717ffc09b8eaed.png" data-download-href="https://discuss.pytorch.org/uploads/default/cc58633c966c2d2cd4419de876717ffc09b8eaed" title="fig1"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/c/cc58633c966c2d2cd4419de876717ffc09b8eaed_2_500x500.png" alt="fig1" data-base62-sha1="t9ILXOez1hbAzUcWuIPlFQBm0cB" width="500" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/c/cc58633c966c2d2cd4419de876717ffc09b8eaed_2_500x500.png, https://discuss.pytorch.org/uploads/default/optimized/2X/c/cc58633c966c2d2cd4419de876717ffc09b8eaed_2_750x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/c/cc58633c966c2d2cd4419de876717ffc09b8eaed_2_1000x1000.png 2x" data-dominant-color="FDFEFE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">fig1</span><span class="informations">1100×1100 15.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: adamr
cooked: <p>Personally I use <code>torch.sigmoid</code> as clamping function. It is more expensive, but the gradients (almost) never vanish.</p>

username: gfox
cooked: <aside class="quote no-group" data-username="always" data-post="3" data-topic="52404">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/5e9695/48.png" class="avatar"> always:</div>
<blockquote>
<p>The only remaining issue is that I do not know how to pass the min/max-values as an argument.</p>
</blockquote>
</aside>
<p>I found <a class="mention" href="/u/always">@always</a>’s  solution quite elegant for my own use case and so have solved the problem of the min/max arguments:</p>
<pre><code class="lang-auto">from torch.cuda.amp import custom_bwd, custom_fwd


class DifferentiableClamp(torch.autograd.Function):
    """
    In the forward pass this operation behaves like torch.clamp.
    But in the backward pass its gradient is 1 everywhere, as if instead of clamp one had used the identity function.
    """

    @staticmethod
    @custom_fwd
    def forward(ctx, input, min, max):
        return input.clamp(min=min, max=max)

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output):
        return grad_output.clone(), None, None


def dclamp(input, min, max):
    """
    Like torch.clamp, but with a constant 1-gradient.
    :param input: The input that is to be clamped.
    :param min: The minimum value of the output.
    :param max: The maximum value of the output.
    """
    return DifferentiableClamp.apply(input, min, max)
</code></pre>