Prepare_qat on module removes hooks

zetyquickly: <p>Hello,</p>
<p>I’ve just started to dive into quantization tools that were introduced in version 1.3.<br>
For now I am trying to train a network with existing model generation code. The code has certain subtleties, one of these are <code>_forward_pre_hooks</code> in several submodules. (See code below)</p>
<p>Here is the problem. After <code>prepare_qat</code> with default config changes submodules to ones with fake quantization and hooks are disappeared. Is it possible to prevent hooks from disappearing during <code>prepare_qat</code> (and <code>submodule.fuse()</code> too)</p>
<p>There is an intermediate code:</p>
<pre><code class="lang-auto">...
print('pre qat: ', model.backbone.bottom_up.blocks[2][3].conv_pwl)
print('pre qat: ', model.backbone.bottom_up.blocks[2][3].conv_pwl._forward_pre_hooks.values())
model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')
torch.quantization.prepare_qat(model, inplace=True)
print('post qat: ', model.backbone.bottom_up.blocks[2][3].conv_pwl)
print('post qat: ', model.backbone.bottom_up.blocks[2][3].conv_pwl._forward_pre_hooks.values())
...
</code></pre>
<p>And the output is:</p>
<pre><code class="lang-auto">pre qat:  Conv2d(120, 40, kernel_size=(1, 1), stride=(1, 1), bias=False)                                                                                                                                    
pre qat:  odict_values([functools.partial(&lt;bound method FeatureHooks._collect_output_hook of &lt;timm.models.feature_hooks.FeatureHooks object at 0x7fd3b5bcf5d0&gt;&gt;, 'blocks.2.3.conv_pwl')])
post qat:  Conv2d(
  120, 40, kernel_size=(1, 1), stride=(1, 1), bias=False
  (activation_post_process): FakeQuantize(
    fake_quant_enabled=True, observer_enabled=True,            scale=tensor([1.]), zero_point=tensor([0])
    (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
  )
  (weight_fake_quant): FakeQuantize(
    fake_quant_enabled=True, observer_enabled=True,            scale=tensor([1.]), zero_point=tensor([0])
    (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
  )
)
post qat:  odict_values([])
</code></pre>

zetyquickly: <p>As I can to prevent hooks from disappearing it is needed to put relevant code somewhere here (during <code>prepare_qat</code>):<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L335" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L335" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/quantization/quantize.py#L335</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="325" style="counter-reset: li-counter 324 ;">
<li>        mod: input module</li>
<li>        mapping: a dictionary that maps from nn module to nnq module</li>
<li>
</li>
<li>    Return:</li>
<li>        The corresponding quantized module of `mod`</li>
<li>    """</li>
<li>    new_mod = mod</li>
<li>    # Always replace dequantstub with dequantize</li>
<li>    if hasattr(mod, 'qconfig') and mod.qconfig is not None or type(mod) == DeQuantStub:</li>
<li>        if type(mod) in mapping:</li>
<li class="selected">            new_mod = mapping[type(mod)].from_float(mod)</li>
<li>    return new_mod</li>
<li>
</li>
<li>def get_observer_dict(mod, target_dict, prefix=""):</li>
<li>    r"""Traverse the modules and save all observers into dict.</li>
<li>    This is mainly used for quantization accuracy debug</li>
<li>    Args:</li>
<li>        mod: the top module we want to save all observers</li>
<li>        prefix: the prefix for the current module</li>
<li>        target_dict: the dictionary used to save all the observers</li>
<li>    """</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>But during fusion it is not that obvious because we map up to three modules to one and each of them could have hooks. I think it is possible to work around with <code>torch.quantization.fuse_modules(...fuser_func=&lt;func&gt;...)</code></p>

jerryzh168: <p>Yeah feel free to submit a PR for preserving the pre hooks in swapping.</p>
<p>For fusion I think we probably need to error out if you have a prehook for intermediate moduels like BatchNorm because when we fuse batchnorm into conv, batchnorm is gone.</p>

zetyquickly: <p>I’ll try!</p>
<p>So Jerry could you explain please for what purpose that was done?</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L99" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L99" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/quantization/quantize.py#L99</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="89" style="counter-reset: li-counter 88 ;">
<li>        if type(child) == nnq.FloatFunctional:</li>
<li>            if hasattr(child, 'qconfig') and child.qconfig is not None:</li>
<li>                child.activation_post_process = child.qconfig.activation()</li>
<li>        else:</li>
<li>            add_observer_(child)</li>
<li>
</li>
<li>    # Insert observers only for leaf nodes, note that this observer is for</li>
<li>    # the output of the module, for input QuantStub will observe them</li>
<li>    if hasattr(module, 'qconfig') and module.qconfig is not None and \</li>
<li>       len(module._modules) == 0 and not isinstance(module, torch.nn.Sequential):</li>
<li class="selected">        # observer and hook will be gone after we swap the module</li>
<li>        module.add_module('activation_post_process', module.qconfig.activation())</li>
<li>        module.register_forward_hook(_observer_forward_hook)</li>
<li>
</li>
<li>def add_quant_dequant(module):</li>
<li>    r"""Wrap the leaf child module in QuantWrapper if it has a valid qconfig</li>
<li>    Note that this function will modify the children of module inplace and it</li>
<li>    can return a new module which wraps the input module as well.</li>
<li>
</li>
<li>    Args:</li>
<li>        module: input module with qconfig attributes for all the leaf modules</li>
</ol></code></pre>


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

<p>After <code>prepare()</code> is called <code>convert()</code> will remove all hooks.<br>
I think there’s a reason to create such hook than to remove it should we preserve all pre forward / post forward hooks except this one?</p>

jerryzh168: <p>Yeah I think we should preserve it, but we need to consider each case carefully, since this is interfering with the quantization. That is, when do we run pre forward hook and post forward hook, do we run it before/after observation/quantization?</p>

zetyquickly: <p>Hello, again <a class="mention" href="/u/jerryzh168">@jerryzh168</a>,</p>
<p>Have a look at this example of pre forward hook run. Here’s EfficientNet implementation that assumes it can be integrated as backbone to FPN (<a href="https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/efficientnet.py#L471" rel="nofollow noopener">https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/efficientnet.py#L471</a>)<br>
So here we can see that during forward blocks are called sequentially than we collect input features from specific layers.</p>
<p>Assume that we have <code>block1</code> who outputs some data, <code>block2</code> who has pre forward hook and directly obtains data from first one and <code>block3</code> who waits for the same data but obtains it using hook of second one.</p>
<p>In this particular example (EfficientNet implementation) during preparartion pre forward hook on <code>block2</code> should be called after observation and therefore after quantization (because we collected statistics for that). If <code>block2</code> happens to be the first in a row who work with quantized data it is very likely that <code>block3</code> works with it too. Anyway we can place dequant before <code>block3</code>.<br>
As for post forward hooks we do not modify input of the module so run them after observation and quantization</p>
<p>Please leave your thoughts when you have time</p>

jerryzh168: <aside class="quote no-group" data-username="zetyquickly" data-post="6" data-topic="72137">
<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/zetyquickly/48/22489_2.png" class="avatar"> zetyquickly:</div>
<blockquote>
<p>Assume that we have <code>block1</code> who outputs some data, <code>block2</code> who has pre forward hook and directly obtains data from first one and <code>block3</code> who waits for the same data but obtains it using hook of second one.</p>
</blockquote>
</aside>
<p>I think this requires that the hooks will do some meaningful computation for both quantized and unquantized data. But as long as we define the semantics clearly it should be OK. So after quantization, the for pre_forward hooks will work with quantized input from previous later, and the forward hooks will work with quantized output from current later, right?</p>

jerryzh168: <aside class="quote no-group" data-username="zetyquickly" data-post="6" data-topic="72137">
<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/zetyquickly/48/22489_2.png" class="avatar"> zetyquickly:</div>
<blockquote>
<p>Assume that we have <code>block1</code> who outputs some data, <code>block2</code> who has pre forward hook and directly obtains data from first one and <code>block3</code> who waits for the same data but obtains it using hook of second one.</p>
</blockquote>
</aside>
<p>actaully we use hooks to implement observe and fake quantize as well, so please make sure that works.</p>

zetyquickly: <p>Yes, I think so.<br>
pre_forward hooks work with quantized data from previous layer, forward hooks work with quantized output from current layer</p>

zetyquickly: <p>You are right, we should check whether we are trying to preserve observer or right hook.<br>
In my <a href="https://github.com/pytorch/pytorch/blob/93540f8bf523e2b365e500549308f321e544b120/torch/quantization/quantize.py#L342" rel="nofollow noopener">PR</a> I have handled that case. Without it provided test set fails, with it works well.</p>
<p>I think I should extended test set to test new functionality. Would you mind to guide me in that?</p>

zetyquickly: <p>Also anticipating possible questions</p>
<p>I’ve introduced changes to <a href="https://github.com/pytorch/pytorch/blob/93540f8bf523e2b365e500549308f321e544b120/torch/quantization/fuse_modules.py#L113" rel="nofollow noopener"><code>fuse_modules.py</code></a></p>
<p>I propose that it is needed to preserve pre and post hooks on fused modules, where it possible. While it hard to define how to preserve hooks for second and third module in sequence (because input data changes and three modules are collapsing into atomic one). But we can easily preserve pre_forward hook of base module.</p>
<p>What cases can we process also?</p>

jerryzh168: <p>Can you post the PR?</p>

zetyquickly: <p>Here it is<br>
</p><aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/37233" target="_blank" rel="nofollow noopener">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/37233" target="_blank" rel="nofollow noopener">Quantization: preserving pre and post forward hooks</a>
    </h4>

    <div class="branches">
      <code>pytorch:master</code> ← <code>zetyquickly:zetyquickly/preserve-hooks</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-04-24" data-time="11:29:54" data-timezone="UTC">11:29AM - 24 Apr 20 UTC</span>
      </div>

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

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

  </div>
</div>

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