[torch::autograd::Function] Any way to know whether forward is called with autograd on?

Hamster: <p>Assume I have a working <code>torch::autograd::Function</code>, and that for the forward method I actually have two ways of computing it, call them <code>forward_1</code> and <code>forward_2</code>. They are such that if one were to do only inference, <code>forward_1</code> would be faster, and if one intends to backprop then <code>forward_2</code>+<code>backward</code> is faster than <code>forward_1</code>+<code>backward</code> (basically <code>forward_2</code> will pre-compute some quantities in anticipation of <code>libtorch</code> later calling <code>backward</code> which can then optimally build upon those quantities to compute the needed gradients efficiently).</p>
<p>So ideally, I’d love to know if there’s a way that could tell me inside <code>torch::autograd::Function::forward</code> whether this method is being called with autograd on or off (ie with or without a no_grad guard), so that I could do something like:</p>
<pre><code class="lang-auto">if(autograd_is_on){
  forward_2(...);
  ctx-&gt;saved_data["..."] = ...; // save the pre-computed stuff that backward will need
} else {
  forward_1(...);
}
</code></pre>

ptrblck: <p>I think <code>torch::GradMode::is_enabled()</code> should work as an internal check.</p>

Hamster: <p>Thank you! I will try this.</p>

Hamster: <p>It doesn’t seem to work :-(. I get a runtime error where torch complains that he’s expecting a Tensor but finds a None instead in ctx-&gt;saved_data in backward during training. When I replace the above pseudo-code with the following the code runs fine:</p>
<pre><code class="lang-auto">forward_2(...);
ctx-&gt;saved_data["..."] = ...; // save the pre-computed stuff that backward will need
</code></pre>
<p>(ie I don’t check whether autograd is on now and I systematically compute what’s needed for backward)</p>
<p>Looks like <code>Gradmode::is_enabled()</code> doesn’t carry the correct value when inside a <code>torch::autograd::Function</code>?</p>

ptrblck: <p>I’m don’t fully understand the current issue, as it seems using the condition raises the runtime error when accessing <code>ctx-&gt;saved_data</code>, but it works without the condition?</p>

Hamster: <p>Assume the code looks like the following:</p>
<pre><code class="lang-auto">torch::autograd::variable_list SomeFunc::forward(torch::autograd::AutogradContext *ctx, torch::Tensor x) {
  /* ... */
  auto stream = c10::cuda::getCurrentCUDAStream();
  if(torch::GradMode::is_enabled()){
    // This will perform the forward computation in a different way using 
    // another slightly slower (in terms of inference only) formula having 
    // subexpressions that are in common with the backward computation, 
    // which makes the forward+backward fast.
    my_cuda_kernels::forward_and_precompute_for_backward(x, out, quantities_for_backward, stream);
    ctx-&gt;saved_data["quantities_for_backward"] = quantities_for_backward;
  } else {
    // The following will take a much faster path,
    // assumes we only need to do inference.
    my_cuda_kernels::forward_only(x, out, stream);
  }          
  return {out};
}

torch::autograd::variable_list SomeFunc::backward(torch::autograd::AutogradContext *ctx, torch::autograd::variable_list grad_output) {
  auto quantities_for_backward = ctx-&gt;saved_data["quantities_for_backward"].toTensor();
  /* do the backward computation by using "quantities_for_backward" to avoid recomputing some subexpressions that are common with the forward computation*/
}
</code></pre>
<p>The problem arises when I use this Function with autograd enabled (ie for example <code>SomeFunc::apply(some_input).mean().backward()</code> assuming <code>some_input</code> has <code>requires_grad</code> set to <code>true</code>), as during runtime libtorch will complain about getting a None instead of a Tensor when arriving at <code>ctx-&gt;saved_data["quantities_for_backward"].toTensor();</code>. So I assume that during the forward phase <code>torch::GradMode::is_enabled()</code> is not returning <code>true</code>.</p>
<p>When I remove the <code>if</code> statement in my forward method and just keep the version that precomputes the backward quantities and saves them in <code>ctx</code> in all cases, the code works again, which further confirms that <code>torch::GradMode::is_enabled()</code> is not working as intended inside a <code>torch::autograd::Function</code>.</p>

ptrblck: <p>Your issue makes sense, since gradient calculation is disabled in custom <code>autograd.Function</code>s by default.<br>
You could check it in <a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html">this example</a> via <code>print(torch.is_grad_enabled())</code> in the <code>forward</code> and would see that it’s returning <code>False</code> so you would have to manually enable it.<br>
Sorry for missing it, but I thought the question is asking how to check if gradient calculation is enabled or not.</p>

Hamster: <p>So I guess then there’s no <code>libtorch</code>-“standard” way I can know inside the <code>forward</code> method whether the <code>autograd::Function</code> is being called from a context with autograd on? (except of course adding a flag myself that I manually set to true when I’m in training/backprop mode and false otherwise)</p>

Hamster: <p>Okay I think I could just encapsulate it in a module and use the <code>is_training</code> method to help me decide which of the two <code>forward</code> versions I choose, this assumes I’ll have to call <code>.train()</code> and <code>.eval()</code> manually but at least it’s more “standard” than introducing my own flags. Probably something similar to BatchNorm.</p>