Custom autograd function with parameters

username: KamikaziZen
cooked: <p>Hello, I am writing a custom function that inherits from <code>torch.autograd.Function</code>. This function uses some parameter(like <code>pdrop</code> for dropout, for example) which is needed to compute gradients.</p>
<pre><code class="lang-auto">class DummyFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, p):
        ctx.save_for_backward(input)
        return p * (5 * input ** 3 - 3 * input)

    @staticmethod
    def backward(ctx, grad_output, p): # not sure I can put p in the inputs
        input, = ctx.saved_tensors
        # number of gradients is required to match the number of inputs
        # even if they are not parameters with "requires_grad" flag
        return grad_output * 3* p * (5 * input ** 2 - 1), None


dummy_function = DummyFunction.apply
</code></pre>
<p>Then I write a module for this function(I assume this is the only way to access p in backward call without saving p in context?)</p>
<pre><code class="lang-auto">class DummyModule(nn.Module):
    def __init__(self, p):
        super().__init__()
        self.p = p

    def forward(self, input):
        return dummy_function(input, self.p)

    def backward(self, grad_outputs):
        return DummyFunction.backward(grad_outputs, self.p) # not sure
</code></pre>
<p>Should I call to <code>DummyFunction.backward(grad_outputs, self.p)</code> here? Seems legit, but I keep wandering if I will break some torch logic here. After all, the documentation pages insist on using <code>DummyFunction.apply</code> for forward call instead of <code>DummyFunction.forward</code>.</p>

username: ptrblck
cooked: <p>Wouldn’t storing <code>p</code> in the <code>ctx</code> via <code>ctx.p = p</code> work (assuming <code>p</code> is not a tensor)?</p>

username: KamikaziZen
cooked: <p>hmm, I didn’t consider it. This could be a solution. Thank you, I will try it.<br>
Is it the way pdrop is accessed during backprop of torch.nn.functional.dropout?  I looked through backward graph and discovered that dropout saves only boolean mask, so it is not clear  where it gets pdrop to normalize this mask with / (1-pdrop) during backward. Does it really just get it from ctx.pdrop ?<br>
What if I have several calls to this dummy_function during forward. Each ctx object will be different so I will not override my p parameter with later calls?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="KamikaziZen" data-post="3" data-topic="177997">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kamikazizen/48/52683_2.png" class="avatar"> KamikaziZen:</div>
<blockquote>
<p>I looked through backward graph and discovered that dropout saves only boolean mask, so it is not clear where it gets pdrop to normalize this mask with / (1-pdrop) during backward.</p>
</blockquote>
</aside>
<p>PyTorch scales with <code>1 / (1 - p)</code> during the training <a href="https://github.com/pytorch/pytorch/blob/fbdb86c1747737c744ad79b5da6bcbd064dc982e/aten/src/ATen/native/Dropout.cpp#L83">here</a> which allows it to avoid the scaling during inference.<br>
This can also be seen in the outputs:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.ones(2, 2)
drop = nn.Dropout(p=0.5)
print(1 / (1 - 0.5))
# 2.0

out = drop(x)
print(out)
# tensor([[0., 2.],
#         [2., 0.]])

drop.eval()
out = drop(x)
print(out)
# tensor([[1., 1.],
#         [1., 1.]])

drop = nn.Dropout(p=0.2)

print(1 / (1 - 0.2))
# 1.25

out = drop(x)
print(out)
# tensor([[1.2500, 1.2500],
#         [0.0000, 0.0000]])

drop.eval()
out = drop(x)
print(out)
# tensor([[1., 1.],
#         [1., 1.]])
</code></pre>
<p>If you want to recompute the <code>mask</code> instead of storing it you would also need to take care of the used seed. Otherwise resampling the <code>mask</code> would generate a new one.</p>
<aside class="quote no-group" data-username="KamikaziZen" data-post="3" data-topic="177997">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kamikazizen/48/52683_2.png" class="avatar"> KamikaziZen:</div>
<blockquote>
<p>What if I have several calls to this dummy_function during forward. Each ctx object will be different so I will not override my p parameter with later calls?</p>
</blockquote>
</aside>
<p>Yes, this should be the case as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">class MyFun(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, p):
        ctx.save_for_backward(input)
        ctx.p = p
        print(f"p in forward {p}")
        return input * p
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        p = ctx.p
        print(f"p in backward {p}")
        return grad_output * input * p, None
    

fun = MyFun.apply

x = torch.randn(2, 2, requires_grad=True)
out1 = fun(x, 0.1)
# p in forward 0.1

out2 = fun(x, 0.5)
# p in forward 0.5

out1.mean().backward()
# p in backward 0.1

out2.mean().backward()
# p in backward 0.5
</code></pre>

username: KamikaziZen
cooked: <p>Thanks,  your answers were very helpful! I was able to solve my problem.<br>
Only one moment is unclear about dropout. The scaling, as you said, is performed during forward which means that to calculate gradients we have to multiply the loss gradient with <code>1/(1-pdrop)</code> after applying mask. Where does that <code>pdrop</code> come from during backward?</p>

username: ptrblck
cooked: <p>I believe the backward formula is implemented in <a href="https://github.com/pytorch/pytorch/blob/c412056921f1e251bd955bb5fd9bd117d5a97ee5/tools/autograd/derivatives.yaml#L583"><code>derivatives.yaml</code></a>:</p>
<pre><code class="lang-auto">- name: native_dropout(Tensor input, float p, bool? train) -&gt; (Tensor, Tensor)
  input: "GradMode::is_enabled() ? infinitely_differentiable_native_dropout_backward(grad, result1, (!train.has_value() || !train.value() ? 1 : (p == 1 ? 0.0 : 1.0 / (1.0 - p)))) : native_dropout_backward(grad, result1, (!train.has_value() || !train.value() ? 1 : (p == 1 ? 0.0 : 1.0 / (1.0 - p))))"
  result0: "(!train.has_value() || train.value()) ? (p == 1 ? 0.0 : 1.0 / (1.0 - p)) * input_t * result1 : input_t"
</code></pre>

username: KamikaziZen
cooked: <p>Thank you for all the answers!</p>