Tensor forward hook

username: sh0416
cooked: <p>Hi,</p>
<p>Is the “register_forward_hook” exist for tensor type?</p>
<p>I couldn’t find it.</p>
<p>Thanks,</p>

username: albanD
cooked: <p>Hi,</p>
<p>There is no such thing no. Why would you want such a hook? If you have the Tensor, you have its value. There is no need for a forward hook.</p>

username: sh0416
cooked: <p>In my use case,</p>
<pre><code class="lang-auto">class M(nn.Module):
  def __init__(self):
    super().__init__()
    self.l = nn.Linear(10, 5)
    self.l2 = nn.Linear(5, 1)

  def forward(self, x):
    x = self.l(x)
    x = torch.exp(x)   # Want to hook this tensor in forward pass
    x = self.l2(x)
    return x
</code></pre>
<p>Actually, I could make forward hook by separating the former part into a single module something like this.</p>
<pre><code class="lang-auto">class FirstLayer(nn.Module):
  def __init__(self):
    super().__init__()
    self.l1 = nn.Linear(10, 5)

  def forward(self, x):
    x = self.l1(x)
    x = torch.exp(x)
    return x

m = FirstLayer()
m.register_forward_hook(lambda x: x)
</code></pre>
<p>However, I want to keep my module single so that I don’t care about other module.</p>
<p>Thanks,</p>

username: albanD
cooked: <p>But this line <code>    x = torch.exp(x)   # Want to hook this tensor in forward pass</code> is actually executed every time the forward is called. So you can put whatever code you want next to that line and it will run during the forward.</p>

username: sh0416
cooked: <p>I think I elaborate on my example more precisely… First, the reason that I create this thread is to make a code more flexible.</p>
<p>I just want to collect some activation for visualization.</p>
<p>Something like this,</p>
<pre><code class="lang-auto">m = Model()
act = dict()
def f(m, i, o):
  act["m"] = o.detach().cpu()

m.layer.register_forward_hook(f)
</code></pre>
<p>If the target activation comes from module, there is no problem because <code>register_forward_hook</code> is in there. However, how to do the same thing for tensor? In my knowledge, this is the best one…</p>
<pre><code class="lang-auto">class Model(nn.Module):
  def __init__(self, act):
    super().__init__()
    self.act = act

  def register_act(self, act):
    self.act = act

  def forward(self, x):
    x = torch.something(x)
    self.act["m"] = x.detach()
    return x
</code></pre>
<p>While it works, I don’t like it <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=9" title=":blush:" class="emoji" alt=":blush:"> because the process of forward depends on <code>self.act</code>. Whenever I use this module, I have to consider the <code>act</code>. I think the hook makes us removing this irritation. I couldn’t figure out the appropriate use case of tensor version <code>register_forward_hook</code> for tensor, which is the reason that lots of specialists don’t implement this functionality.</p>
<p>Thanks anyway, your advice give me lots of useful insight.</p>

username: albanD
cooked: <p>My point is more that you cannot do such an API. By the time to have the Tensor so that you could register a hook on it, it’s content is available. So there is no need to hook anymore.<br>
And if you want to place a hook beforehand, you have no Tensor to place it onto because the Tensor does not exist yet.</p>

username: ptrblck
cooked: <p>As <a class="mention" href="/u/alband">@albanD</a> explained, you cannot register a hook on a tensor, as it’s already available.<br>
Based on your previous code snippet I guess you could use an <code>nn.Identity</code> layer, pass <code>x</code> to it after the <code>exp</code> operation and register a hook on this “fake” layer. I don’t think it would be cleaner than storing the tensor as in your last code snippet, but you might prefer this approach.</p>

username: sh0416
cooked: <p>That is the fresh thought, I didn’t aware of the existence of <code>nn.Identity</code>. <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=9" title=":blush:" class="emoji" alt=":blush:"></p>
<p>Maybe the code will be like this… (This is for my future. I usually revisit my thread to copy some snippet)</p>
<pre><code class="lang-auto">class Model(nn.Module);
  def __init__(self):
    super().__init__()
    self.w = nn.Linear(10, 5)
    self.tracker_act = nn.Identity("exp_after_linear")

  def forward(self, x):
    x = self.w(x)
    x = self.tracker_act(torch.exp(x))
    return x

act = {}
def collect(m, i, o):
  return act["exp_after_linear"] = o.detach().cpu()

model = Model()
model.tracker_act.register_forward_hook(collect)
</code></pre>
<p>I think it is nice. <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=9" title=":blush:" class="emoji" alt=":blush:"><br>
Thanks for the nice discussion, <a class="mention" href="/u/alband">@albanD</a> and <a class="mention" href="/u/ptrblck">@ptrblck</a>.</p>

username: MichalM
cooked: <p>You can also implement a wrapper <code>class</code> for a function.<br>
This also allows you to keep track of connections between modules.<br>
Calling Identity after <code>torch.exp</code> returns (via hook) two identical tensors – <code>*input_tensors, output_tensor</code>.<br>
The connection to the previous module, by sharing the same tensor of previous module output and input of considered module, is lost.<br>
Using a wrapper module allows for a continuous connection.<br>
*Until the previous tensors are released from memory and the new ones get identical ptrs.</p>
<pre><code class="lang-auto">import torch


class Wrapper(torch.nn.Module):
    def __init__(self, operation, description:str=""):
        super().__init__()
        self.operation = operation
        self.description = description

    def extra_repr(self,):
        return f"operation={self.description}, description=\"{self.description}\""

    def forward(self, *x):
        return self.operation(*x)


class Example(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l = torch.nn.Linear(10, 5)
        self.exp = Wrapper(torch.exp,"torch.exp")
        self.l2 = torch.nn.Linear(5, 1)
        self.square = Wrapper(torch.square,"torch.square")

    def forward(self, x):
        x = self.l(x)
        x = self.exp(x)
        x = self.l2(x)
        x = self.square(x)
        return x


def hook(module,in_tensors,out_tensor):
    print(module,[t.data_ptr() for t in in_tensors], out_tensor.data_ptr())



if __name__ == '__main__':
    m = Example()
    
    hooks = []
    for ch in m.children():
        hooks.append(ch.register_forward_hook(hook)
                     )
    t = torch.rand(1,10)
    print("Input t:", t.data_ptr())
    o = m(t)
    print("Output o:", o.data_ptr())
    
    for h in hooks:
        h.remove()
</code></pre>
<pre><code class="lang-auto">Input t: 102179200
Linear(in_features=10, out_features=5, bias=True) [102179200] 102739072
Wrapper(operation=torch.exp, description="torch.exp") [102739072] 103698048
Linear(in_features=5, out_features=1, bias=True) [103698048] 103786048
Wrapper(operation=torch.square, description="torch.square") [103786048] 104697856
Output o: 104697856
</code></pre>