Question about prepare_fx

username: cokespace2
cooked: <p>In practice, if I have a <code>nn.Conv2d</code> after a <code>nn.Conv2d</code>, the <code>prepare_fx</code> method would insert an observer between these two modules.</p>
<p>But, if I have a <code>nn.BatchNorm2d</code> after a <code>nn.Conv2d</code> (<strong>NO FUSING</strong>), the <code>prepare_fx</code> would not insert an observer, which makes <code>convert_fx</code> wouldn’t convert the referenced conv2d module to quantized conv2d module.</p>

username: supriyar
cooked: <p>hi, can you share a simple repro here exhibiting the problem? how are you disabling the fusion of Conv+BN in the prepare step?</p>

username: cokespace2
cooked: <p>Sure, basically I just comment the fusing part in <code>prepare_fx</code>, but you can reprod it like below:</p>
<pre><code class="lang-python">class MyModule(nn.Module):
  def __init__(self):
    self.conv1 = nn.Conv2d(3,16,3,1,1)
    self.bn1 = nn.BatchNorm2d(16)
    self.relu = nn.ReLU()
    self.check = nn.Conv2d(16, 32, 3, 1, 1)

  def forward(self, x):
    x = self.check(self.relu(self.bn1(self.conv1(x))))
</code></pre>
<p>and use:</p>
<pre><code class="lang-auto">model = MyModule()
qconfig_mapping = QConfigMapping().set_global(get_default_qconfig("x86"))
prepared = prepare_fx(model, qconfig_mapping)
</code></pre>
<p>The <code>conv1</code> → <code>bn1</code> → <code>relu</code> would be fused to a single <code>conv2d</code> module, and you would see that one observer is inserted between this <code>conv2d</code> and <code>check</code> module.<br>
But, if we replace <code>self.check</code> to <code>nn.BatchNorm2d(16)</code>, the observer would not be inserted.</p>

username: supriyar
cooked: <p>When I run this snippet locally, I see the observers inserted after conv1 and before <code>self.check</code>.<br>
My test code</p>
<pre><code class="lang-auto">      class M(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv1 = nn.Conv2d(3,16,3,1,1)
                self.bn1 = nn.BatchNorm2d(16)
                self.relu = nn.ReLU()
                self.check = nn.BatchNorm2d(16) #nn.Conv2d(16, 32, 3, 1, 1)

            def forward(self, x):
                x = self.check(self.relu(self.bn1(self.conv1(x))))
                return x

        model = M().eval()
        mp = prepare_fx(model, get_default_qconfig_mapping(), example_inputs=torch.randn(1, 3, 1, 1))
        print(mp)
        mc = convert_fx(mp)
        print(mc)
</code></pre>
<p>Here is the output from prepare</p>
<pre><code class="lang-auto">GraphModule(
  (activation_post_process_0): HistogramObserver(min_val=inf, max_val=-inf)
  (conv1): ConvReLU2d(
    (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
  )
  (activation_post_process_1): HistogramObserver(min_val=inf, max_val=-inf)
  (check): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (activation_post_process_2): HistogramObserver(min_val=inf, max_val=-inf)
)



def forward(self, x):
    activation_post_process_0 = self.activation_post_process_0(x);  x = None
    conv1 = self.conv1(activation_post_process_0);  activation_post_process_0 = None
    activation_post_process_1 = self.activation_post_process_1(conv1);  conv1 = None
    check = self.check(activation_post_process_1);  activation_post_process_1 = None
    activation_post_process_2 = self.activation_post_process_2(check);  check = None
    return activation_post_process_2
</code></pre>
<p>output after convert</p>
<pre><code class="lang-auto">GraphModule(
  (conv1): QuantizedConvReLU2d(3, 16, kernel_size=(3, 3), stride=(1, 1), scale=1.0, zero_point=0, padding=(1, 1))
  (check): QuantizedBatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)



def forward(self, x):
    conv1_input_scale_0 = self.conv1_input_scale_0
    conv1_input_zero_point_0 = self.conv1_input_zero_point_0
    quantize_per_tensor = torch.quantize_per_tensor(x, conv1_input_scale_0, conv1_input_zero_point_0, torch.quint8);  x = conv1_input_scale_0 = conv1_input_zero_point_0 = None
    conv1 = self.conv1(quantize_per_tensor);  quantize_per_tensor = None
    check = self.check(conv1);  conv1 = None
    dequantize_2 = check.dequantize();  check = None
    return dequantize_2
</code></pre>

username: cokespace2
cooked: <p>My bad, I shouldn’t add <code>self.relu</code>, so here’s what I got:</p>
<pre><code class="lang-auto">class M(nn.Module):
  def __init__(self):
    super().__init__()
    self.conv1 = nn.Conv2d(3,16,3,1,1)
    self.bn1 = nn.BatchNorm2d(16)
    self.check = nn.BatchNorm2d(16) #nn.Conv2d(16, 32, 3, 1, 1)

  def forward(self, x):
    x = self.check((self.bn1(self.conv1(x)))
    return x
</code></pre>
<p>and after <code>prepare</code>:</p>
<pre><code class="lang-auto">GraphModule(
  (activation_post_process_0): HistogramObserver(min_val=inf, max_val=-inf)
  (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (check): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (activation_post_process_1): HistogramObserver(min_val=inf, max_val=-inf)
)
</code></pre>
<p>after <code>convert</code>:</p>
<pre><code class="lang-auto">GraphModule(
  (conv1): QuantizedConv2d(Reference)(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (check): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
</code></pre>

username: 111357
cooked: <p>I think the <code>self.conv1</code> already converted into <code>QuantizedConv2d</code>.</p>

username: cokespace2
cooked: <p><code>conv1</code> is still a <strong>Reference</strong> convolution which not takes uint8 dtype tensor input, and <code>check</code> is still a float point version of <code>BatchNorm2d</code>.</p>
<p>I’d like to have a <code>QuantizedConv2d</code> and a <code>QuantizedBatchNorm2d</code> module like <a class="mention" href="/u/supriyar">@supriyar</a> comment above.</p>

username: supriyar
cooked: <p>this looks like a bug, can you file an issue on <a href="https://github.com/pytorch/pytorch/issues" class="inline-onebox" rel="noopener nofollow ugc">Issues · pytorch/pytorch · GitHub</a> with the quantization label? Thanks!</p>

username: cokespace2
cooked: <p>Sure, I’ll open one, btw is there any way that we can walkaround this potential bug?</p>

username: HDCharles
cooked: <p>This isn’t working because you’re using default settings for something that is far from default. In general, Conv + multiple Batchnorms should be the same as Conv + 1 Batchnorm which is matematically equivalent to just a Conv op with a different weight and thats what we’re handling by default, fusing conv-bn → conv. The issue is its trying to do like a fused quantization with conv-bn but the code assumes that they’ve been fused together and that there’s only a conv but because of the multiple batchnorms, only 1 gets fused and the rest cause a conflict.</p>
<p>If you really needed to get around this the easiest way is that you could either insert a torch.nn.Identity between the 2 batchnorms in your model or you could make a new backend_config that works with the above situation, which would be equivalent to just commenting out these lines from the default backend config</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/3d6fd07c46e23250c7e12b6965b4e8b0f85df4c1/torch/ao/quantization/backend_config/_common_operator_config_utils.py#L305-L309">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/3d6fd07c46e23250c7e12b6965b4e8b0f85df4c1/torch/ao/quantization/backend_config/_common_operator_config_utils.py#L305-L309" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/3d6fd07c46e23250c7e12b6965b4e8b0f85df4c1/torch/ao/quantization/backend_config/_common_operator_config_utils.py#L305-L309" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/3d6fd07c46e23250c7e12b6965b4e8b0f85df4c1/torch/ao/quantization/backend_config/_common_operator_config_utils.py#L305-L309</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="305" style="counter-reset: li-counter 304 ;">
          <li>conv_configs.append(</li>
          <li>    BackendPatternConfig((convs.root, convs.bn))</li>
          <li>        .set_dtype_configs(dtype_configs)  # noqa: E131</li>
          <li>        .set_fuser_method(fuse_conv_bn)</li>
          <li>        .set_fused_module(convs.fused_conv_bn))</li>
      </ol>
    </code></pre>



  </article>

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

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