Defining additional methods in an nn.Module

tjoseph: <p>Hi everyone,<br>
I am wondering whether it is ok to define additional function in an nn.Module:</p>
<pre><code class="lang-auto">class MyModule(nn.Module):
    ...
    def classify(...):
        ...
    def detect(...):
        ....
    def forward(...):
        ....
</code></pre>
<p>My question is based on this previous thread: <a href="https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690" class="inline-onebox">Any different between model(input) and model.forward(input)</a> which says that one should not call forward directly because some hooks are not used then. So is it discouraged define additional methods? Putting everything in <code>forward()</code> with some flag would be an alternative, but an ugly one.</p>

InnovArul: <p>I am not sure about your exact use case. In my understanding, if you are not using any <code>forward</code> or <code>backward</code> hooks, you can always define additional methods inside <code>nn.Module</code>.</p>
<p>Note that the hooks would only work if <code>model(input)</code> is called, which in turn processes the hooks internally before and after calling <code>forward()</code>.</p>
<p>You can have a look at the code here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/6f1455440b1eae1195b1257abd412479c7873499/torch/nn/modules/module.py#L1050-L1100">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/6f1455440b1eae1195b1257abd412479c7873499/torch/nn/modules/module.py#L1050-L1100" target="_blank" rel="noopener">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/6f1455440b1eae1195b1257abd412479c7873499/torch/nn/modules/module.py#L1050-L1100" target="_blank" rel="noopener">pytorch/pytorch/blob/6f1455440b1eae1195b1257abd412479c7873499/torch/nn/modules/module.py#L1050-L1100</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="1050" style="counter-reset: li-counter 1049 ;">
          <li>def _call_impl(self, *input, **kwargs):</li>
          <li>    forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.forward)</li>
          <li>    # If we don't have any hooks, we want to skip the rest of the logic in</li>
          <li>    # this function, and just call forward.</li>
          <li>    if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks</li>
          <li>            or _global_forward_hooks or _global_forward_pre_hooks):</li>
          <li>        return forward_call(*input, **kwargs)</li>
          <li>    # Do not call functions when jit is used</li>
          <li>    full_backward_hooks, non_full_backward_hooks = [], []</li>
          <li>    if self._backward_hooks or _global_backward_hooks:</li>
          <li>        full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()</li>
          <li>    if _global_forward_pre_hooks or self._forward_pre_hooks:</li>
          <li>        for hook in (*_global_forward_pre_hooks.values(), *self._forward_pre_hooks.values()):</li>
          <li>            result = hook(self, input)</li>
          <li>            if result is not None:</li>
          <li>                if not isinstance(result, tuple):</li>
          <li>                    result = (result,)</li>
          <li>                input = result</li>
          <li>
          </li>
<li>    bw_hook = None</li>
      </ol>
    </code>
  </pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/6f1455440b1eae1195b1257abd412479c7873499/torch/nn/modules/module.py#L1050-L1100" target="_blank" rel="noopener">show original</a>

  </article>

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

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


tjoseph: <p>Thanks for your answer.</p>
<p>I guess this means there should be no problems to use other methods for inference and <code>model()</code> for training.</p>

InnovArul: <p>Yes you are correct.</p>