PyTorch forward hook cannot capture all input variables

zjplab: <p>I have a program with a sub neural network whose <code>forward</code> method takes three input arguments:</p>
<pre><code class="lang-auto">def forward(self, input, hidden,
                transformed_input=None):
</code></pre>
<p>But in reality it was fed with the latter two arguments, namely <code>hidden</code>  and <code>tranformed_input</code>.<br>
I registered one forward hook and found only two inputs have been captureds:</p>
<pre><code class="lang-auto">def hook(module, input, output):
    print(len(input)) #it outputs 2
</code></pre>
<p>When I am trying to print <code>input</code> in my forward hook:</p>
<pre><code class="lang-auto">def hook(module, input, output):
    #print(module)
    print(input)
</code></pre>
<p>It outputs a tuple of size 2 with (None, (some tensor, some tensor)).</p>
<p>Notice in my program <code>hidden</code> does have this kind of (some tensor, some tensor) shape. My best guess so far is that <strong>PyTorch ignores some parts whose default value is None</strong>. Since in my program in the forward pass i<code>nput</code> was fed with None and <code>transformed_input</code> 's default vallue is None, hence causing the above observation.</p>
<p>However, I did another experiement:</p>
<pre><code class="lang-auto">class Test(nn.Module):
  def __init(self):
    pass
  def forward(self, input1, input2, input3=None):
    if input1 is not None:
      return input1+input2
    else:
      return input2

a=Test() 

def hook(module, input, output):
  print(module)
  print(len(input))
  print(input)

a.register_forward_hook(hook)
a(None, torch.tensor([1]), torch.tensor([9]))
# Test()
# 3
# (None, tensor([1]), tensor([9]))
#tensor([1]) 

</code></pre>
<p>which correctly output everything. So what are your thoughts on this?</p>

ptrblck: <p>My guess it you might be registering or looking at the wrong module in your initial code, as the small code snippet seems to work.</p>
<p>Try to add a name to the hook so that the print statement will show you which layer was called with which inputs and outputs.</p>

zjplab: <p>I think I found the bug.<br>
My new hook looks like this:</p>
<pre><code class="lang-auto">def hook(module, input, output):
    print(module)
    print(len(input))
    print(input)
</code></pre>
<p>When in the forward contains the explicit variable name:</p>
<pre><code class="lang-auto">cell(None, (torch.zeros(1,1150), torch.zeros(1,115,10)), transformed_input=cell.ih(torch.ones(1,400)))
</code></pre>
<p>It outputs:</p>
<blockquote>
<p>ONLSTMCell(<br>
(ih): Sequential(<br>
(0): Linear(in_features=400, out_features=4830, bias=True)<br>
)<br>
(hh): LinearDropConnect(in_features=1150, out_features=4830, bias=True)<br>
)<br>
2<br>
(None, (tensor([[0., 0., 0.,  …, 0., 0., 0.]]), tensor([[[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
…,<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.]]])))</p>
</blockquote>
<p>When variable name in forward method is dropped:</p>
<pre><code class="lang-auto">cell(None, (torch.zeros(1,1150), torch.zeros(1,115,10)), cell.ih(torch.ones(1,400)))
</code></pre>
<blockquote>
<p>ONLSTMCell(<br>
(ih): Sequential(<br>
(0): Linear(in_features=400, out_features=4830, bias=True)<br>
)<br>
(hh): LinearDropConnect(in_features=1150, out_features=4830, bias=True)<br>
)<br>
3<br>
(None, (tensor([[0., 0., 0.,  …, 0., 0., 0.]]), tensor([[[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
…,<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.],<br>
[0., 0., 0.,  …, 0., 0., 0.]]])), tensor([[ 0.2503,  0.5609, -0.3321,  …,  0.2259, -0.2337, -0.3209]],<br>
grad_fn=))</p>
</blockquote>

zjplab: <p>And I have already issued this bug on PyTorch’s github repo</p>

jspaezp: <p>I was able to generate an example that reproduces the problem and after some digging …</p>
<p>It is a reflection of the problems discussed here:</p>
<p>Issues In the pytorch github:</p>
<ul>
<li>59923</li>
<li>35643</li>
</ul>
<p>and it seems like a fix for kwargs to be handled correctly by the hooks is in the making</p>
<aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/pytorch/pull/61606">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/pull/61606" 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/61606" target="_blank" rel="noopener nofollow ugc">[pytorch] Lazy Module Fix for kwargs argument support on forward</a>
    </h4>

    <div class="branches">
      <code>pytorch:master</code> ← <code>ertrue:export-D29687148</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-07-13" data-time="22:23:24" data-timezone="UTC">10:23PM - 13 Jul 21 UTC</span>
      </div>

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

      <div class="lines" title="1 commits changed 3 files with 59 additions and 15 deletions">
        <a href="https://github.com/pytorch/pytorch/pull/61606/files" target="_blank" rel="noopener nofollow ugc">
          <span class="added">+59</span>
          <span class="removed">-15</span>
        </a>
      </div>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Summary:
This is a fix to resolve the issue: #35643 and #59923

Test Plan:
p<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/pull/61606" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">ython3 test: caffe2/test/test_nn.py

Differential Revision: D29687148</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>TLDR: Keywords argument dont show up in hooks, fix your code until they do</p>
<pre><code class="lang-auto">import torch
from torch import nn


class BorinModel(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.linear1 = nn.Linear(in_features=10, out_features=1)
        self.linear2 = nn.Linear(in_features=20, out_features=1)

    def forward(self, x, y, debug=True):
        if debug:
            print("im debuggin")

        return self.linear1(x) + self.linear2(y)

class EvenMoreBoringModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.boring_stuff = BorinModel()

    def forward(self, x, y):
        return self.boring_stuff(x, y, debug=True)

    def forward2(self, x, y):
        return self.boring_stuff(x=x, y=y, debug=True)

def hook(m, i, o):
    print(m)
    print(f"Inputs: {[x.shape for x in i]}")
    print(f"Outputs: {o.shape}")

boring_instance = EvenMoreBoringModel()

handle = boring_instance.boring_stuff.register_forward_hook(hook)

boring_instance.forward(torch.ones(5,10), torch.ones(5, 20))
# im debuggin
# BorinModel(
#   (linear1): Linear(in_features=10, out_features=1, bias=True)
#   (linear2): Linear(in_features=20, out_features=1, bias=True)
# )
# Inputs: [torch.Size([5, 10]), torch.Size([5, 20])]
# Outputs: torch.Size([5, 1])
boring_instance.forward2(torch.ones(5,10), y = torch.ones(5, 20))
# im debuggin
# BorinModel(
#   (linear1): Linear(in_features=10, out_features=1, bias=True)
#   (linear2): Linear(in_features=20, out_features=1, bias=True)
# )
# Inputs: []
# Outputs: torch.Size([5, 1])
</code></pre>