Implementing elegant KFAC - need clarification of hooks

botev: <p>So I’m trying to implement a nice version of KFAC in pytorch which to require only a single call to a function which to augments a <code>nn.Module</code>. I think the hooks mechanism actually can be quite an elegant solution for this, however, I need some clarification, to make sure I understand things correctly.<br>
For example, let’s consider <code>nn.Linear</code>. During the backward the backward hook <code>grad_input</code> is a tuple of 3 variables. It is unclear to me why they are 3 and did not find any docs on what they are (one seems to be the grad with respect to the weights, and one the grad with respect to the inputs, the third most likely is the grad with respect to the <strong>broadcasted</strong> bias):</p>
<pre><code class="lang-auto">weight : torch.Size([784, 128])
bias : torch.Size([784])
d1: torch.Size([64, 784]) # most likely grad of broadcasted bias
d2: torch.Size([64, 128]) # most likely grad of inputs
d3: torch.Size([128, 784]) # most likely grad of weight.T
grad_output: torch.Size([64, 784])
</code></pre>
<p>Why the <code>input</code> to the module is not part of this and is there a way to access it? Generally, modifying gradients would require knowing the input.</p>
<p>Additionally, is there a way to make the Modules produce twice more gradients than then the number of their inputs, but somehow intersect and do separate things with half of them and the other half being normal gradients?</p>

tom: <p>The main clarification of module backward hooks is <a href="https://github.com/pytorch/pytorch/issues/598" rel="nofollow noopener">“don’t use them”</a>.<br>
Sorry, but there isn’t a better explanation.<br>
If you look at <a href="https://github.com/yaroslavvb/kfac_pytorch/blob/master/kfac_pytorch.py" rel="nofollow noopener">Yaroslavvb’s implementation</a>, you see that he wrapped the actual autograd functions in order to have access to all he needs, quite likely, something like that is necessary.</p>
<p>Best regards</p>
<p>Thomas</p>

botev: <p>Ah, that would be quite unfortunate than is this means there is no easy way to apply it to already existing code as it requires that you explicitly use the overridden functions.<br>
Nevertheless, thanks at least that makes it clear what needs to be done.</p>

AlphaBetaGamma96: <p>Hi,</p>
<p>Just to ask a quick question. So, the problem with using <code>register_backward_hook</code> is that it doesn’t always return the correct number of gradients? And, in order to fix this (so that <code>register_backward_hook</code> works as intended) is to explicitly write out the <code>forward</code> and <code>backward</code> pass within a custom <code>torch.autograd.function</code> for all layers and matmul within my code so that the correct number of gradients are returned?</p>
<p>Thank you!</p>

tom: <p>I think you want to move on to <code>register_full_backward_hook</code> on nightlies:</p><aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/46163" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Pull Request">
    <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/pull/46163" target="_blank" rel="noopener nofollow ugc">Fix Module backward hooks for all Tensor inputs/outputs</a>
    </h4>

    <div class="branches">
      <code>pytorch:master</code> ← <code>albanD:bw_hook</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-10-11" data-time="16:26:56" data-timezone="UTC">04:26PM - 11 Oct 20 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/albanD" target="_blank" rel="noopener nofollow ugc">
          <img alt="albanD" src="https://avatars.githubusercontent.com/u/6359743?v=4" class="onebox-avatar-inline" width="20" height="20">
          albanD
        </a>
      </div>

      <div class="lines" title="10 commits changed 6 files with 491 additions and 50 deletions">
        <a href="https://github.com/pytorch/pytorch/pull/46163/files" target="_blank" rel="noopener nofollow ugc">
          <span class="added">+491</span>
          <span class="removed">-50</span>
        </a>
      </div>
    </div>

  </div>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


AlphaBetaGamma96: <p>Thank you for the response! So, I assume the only way to do KFAC via registering hooks is to use the nightly build? (I’m currently using 1.7 for reference)</p>
<p>Thank you!</p>

tom: <p>You could always build your own, that’s all I know about.<br>
But so I heard rumors that we can expect the 1.8 RC in February, it might not be too far to a release that has what you need.</p>

AlphaBetaGamma96: <p>Great! I’ll look forward to it! Thanks!</p>