Stop model evaluation with forward hook

username: antspy
cooked: <p>Let’s say I have a model that I only want to run until the 8-th layer. I register a forward hook on that layer to save the output, and all works well. The problem is that the network will still run until completion, which makes things slower than necessary.</p>
<p>Is there a way to stopping it early? Or should I raise an exception in the forward hook and catch it outside, to get out of the network forward function?</p>

username: richard
cooked: <p>Do you have any control over the architecture of the network? In the python script for that network I imagine you could just delete the lines after the layer whose output you want.</p>

username: antspy
cooked: <p>Sure, but I want to do this automatically, without manually deleting lines of code every time <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: jpeg729
cooked: <p>You could add an argument to the <code>forward</code> method of your model.</p>
<pre><code class="lang-auto">def forward(self, input, num_layers):
    # automatically use only num_layers</code></pre>

username: antspy
cooked: <p>Yes, sure. But then I would need to manually modify the forward method of every class I am interested in. What I want is to have a layer of abstraction; I don’t care what module it is or how it’s defined, I want to stop execution after a certain forward hook is triggered.</p>

username: richard
cooked: <p>One way to do this is to subclass <a href="http://pytorch.org/docs/master/_modules/torch/nn/modules/container.html#Sequential"><code>nn.Sequential</code></a> and change the forward method so that you can use a number to determine how many layers deep it goes. Then, you can build your models using your modified version of <code>nn.Sequential</code> to do this.</p>

username: antspy
cooked: <p>Again, thanks for the suggestion but this is not what I am looking for. I want to take in input a model, defined by someone else in some way that I don’t care about, and stop execution at the n-th module, as returned by model.modules(). I don’t want to modify the model definition, or do anything that would imply having knowledge of how the model is defined or how it works.</p>
<p>In short I want it to be completely automated and work on any possible model without any kind of modification on my part <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: SimonW
cooked: <p>You need to define more specifically what you want. When you say a “module”, do you mean things in <code>nn.*</code> like linear layers <code>nn.Linear</code>, conv layers <code>nn.Conv*d</code>, upsampling <code>nn.Upsample</code>, nonlinearities, etc.? Or do you mean PyTorch <code>Function</code>s, which are the real fundamental ops of the computation graph, e.g. <code>view</code>, <code>thnn_convolution</code>, <code>svd</code>, <code>transpose</code>, <code>matmul</code>, etc. Each of the former class is usually implemented with one or more <code>Function</code>s from the latter class, so you need to be clear what your notion of n-th “module” means. If you are referring to the former class, which is probably more reasonable because they are higher-level and more intuitive, there is another difficulty. People can write models using lower-level <code>Function</code>s, e.g. instead of using <code>nn.Linear</code> they use a transpose + <code>addmm</code>, use functions, such as linalg ops including <code>trtrs</code> and <code>potrs</code>, which do not have the corresponding notion of a higher level <code>nn.*</code> module.</p>

username: antspy
cooked: <p>As I mentioned before, I will consider a “module” everything that is returned by <code>model.modules()</code>. They are also the only thing you can attach forward hooks to.</p>
<p>Some people may decide to write their models in a different way, such that <code>model.modules()</code> will be empty or missing pieces, as you point out. But since there is no way around it, I will accept this limitation.</p>
<p>Essentially I want this:</p>
<pre><code class="lang-auto">for idx, md in model.named_modules(): 
    if idx == idx_to_stop_execution_to:
        md.register_forward_hook(my_function)
_ = model(inputs)
</code></pre>
<p>and <code>my_function</code> is defined in such a way that it will save the output of that module, and then <em>stop the forward pass of model</em> to save time. I tried to raise an exception in my_function, but the outputs were not being correctly saved.</p>
<p>*** EDIT***</p>
<p>It turns out, the error was that I <em>forgot</em> to de-register the forward hooks. The code below works as intended.</p>
<pre><code class="lang-auto">for idx, md in model.named_modules(): 
    if idx == idx_to_stop_execution_to:
        md.register_forward_hook(my_function)
try: _ = model(inputs)
except CustomException: pass
</code></pre>

username: SimonW
cooked: <aside class="quote no-group" data-username="antspy" data-post="9" data-topic="12882">
<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/antspy/48/3475_2.png" class="avatar"> antspy:</div>
<blockquote>
<p>As I mentioned before, I will consider a “module” everything that is returned by model.modules().</p>
</blockquote>
</aside>
<p>This is still ill-defined. What if a module contains <code>Sequential(linear1, linear2)</code>? All three will show up in <code>model.modules()</code>, but one of them contains the other two. What do you mean by n-th module?</p>
<aside class="quote no-group" data-username="antspy" data-post="9" data-topic="12882">
<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/antspy/48/3475_2.png" class="avatar"> antspy:</div>
<blockquote>
<p>It turns out, the error was that I forgot to de-register the forward hooks. The code below works as intended.</p>
<p>for idx, md in model.named_modules():<br>
if idx == idx_to_stop_execution_to:<br>
md.register_forward_hook(my_function)<br>
try: _ = model(inputs)<br>
except CustomException: pass</p>
</blockquote>
</aside>
<p>This won’t work for a bunch of models, e.g.</p>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, x):
  for _ in range(100):
    x = self.fc(x)
  return x
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, x):
  if x.data[0] &lt; 1:
    return self.fc1(x)
  else:
    return self.fc2(x)
  return x
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, x):
  z = self.subnet(Variable(torch.randn(1, 2))
  return z + x
</code></pre>
<p>No where guarantees that <code>named_modules</code> or <code>modules</code> returns in certain order. It doesn’t even work for sequential modules.</p>
<p>If you want to limit youself to <code>nn.Sequential</code>, fine, there are a number of ways to do this. But if you want to do it for general modules. I don’t see how any of the proposals you made could work. You should think about what exactly is a “module”,  and with dynamic python flow control, how to get the n-th actually executed “module”.</p>

username: antspy
cooked: <p>You make some very good points; especially about sequential and the possibility of each module showing up twice.</p>
<p>Note though that</p>
<p>1 - My main question was  “stop execution when forward hook is triggered”. So essentially everywhere you want to put a forward hook in, stop execution there. The idea of having it stopped at the “n-th” module was just for my convenience, so that I don’t have to inspect any piece of code.</p>
<p>2 - While there are significant pitfalls to be aware of (and I thank you for pointing them out!), these are left to the “user”. You put in input a model, and a index n; you will get in output the result of registering a forward hook to the “n-th” module returned by model.modules().  It is up to the user then to identify which index is correct, whether a layer shows up twice, or whether it just doesn’t make any sense (as in your last examples). This is just a utility that is useful in a large enough number of cases, and for the cases where it does not make sense, then it does not make sense. The only thing that is required for this to work is that model.modules() is deterministic; i.e. given the same model definition, the modules are returned in the same order. If this holds, then the user would put the correct n, and the method will work (if applicable).</p>
<p>What do you think? Thanks for the answer anyhow <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: SimonW
cooked: <p>Given the limitations, it might not be very helpful in many cases. But a way to do this may be registering a hook on every submodule, where this hook increments a global counter and if the counter reaches N saves the result at a global ptr and throws an exception. It’s very very hacky though <img src="//discuss.pytorch.org/images/emoji/apple/smiley.png?v=5" title=":smiley:" class="emoji" alt=":smiley:"></p>

username: shivamsaboo17
cooked: <p><a class="mention" href="/u/simonw">@SimonW</a> I am implementing  ManifoldMixup (<a href="http://proceedings.mlr.press/v97/verma19a/verma19a.pdf" rel="nofollow noopener">http://proceedings.mlr.press/v97/verma19a/verma19a.pdf</a>) which is similar to Mixup regularization technique except it works at layer level instead of input level as in Mixup. I need a similar functionality to implement this,</p>
<ol>
<li>Select a random index and apply forward hook to that layer</li>
<li>Forward pass using data input x_0 and record output at hooked layer</li>
<li>Use this output along with new input x_1 by adding new hook at the same layer to do this mixup operation</li>
</ol>
<p>It would be nice if I can stop processing model’s forward pass at hook in first step. Is there a better way to do this now since it has been over a year after this thread?</p>

username: Scott_Maddox
cooked: <p>I’m a bit late to discussion, but to implement ManifoldMixup, rather than stopping the processing after layer K, you could register a forward_hook for layer K that computes the weighted-averages across your mixup dimension and returns the resulting tensor for further processing.</p>