Questions about QAT

smivv: <p>Hello!</p>
<p>I am trying to train MobileNetV3 with Lite Reduced ASPP for Semantic Segmentation using Quantization Aware Training, but for some reason it does not training at all. Output of the model seems to be like random noise.<br>
So I have couple of questions.</p>
<ol>
<li>Currently I have such activations as <code>nn.ReLU6</code>, <code>nn.Sigmoid</code>, <code>nn.Hardsigmoid</code> and <code>nn.Hardswish</code>. I tried both approaches - wrap them with <code>QuantWrapper</code> or replace with ReLU. Both didn’t help. What is the correct way?</li>
<li>I’ve replaced <code>functional.interpolate</code> with <code>nn.UpsamplingBilinear2d</code>. It also didn’t help. But is it relevant or I can use functional analogue?</li>
<li>I’ve replaced all <code>add</code> and <code>mul</code> operations to separated torch.nn.quantized.FloatFunctional(). Didn’t work.</li>
<li>Usually with QAT approach does model training longer or approximately it should take the same time?</li>
<li>Am I right that is possible to train model with QAT using GPU?</li>
</ol>
<p>Here is the print of my model:</p>
<pre><code class="lang-auto">QuantizableMobileNetV3SmallLRASPP(
  (_quant): QuantStub()
  (_encoder): QuantizableMobileNetV3(
    (_layers): ModuleList(
      (0): Sequential(
        (0): ConvBn2d(
          (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (1): Identity()
        (2): Hardswish()
      )
      (1): QuantizableInvertedResidual(
        (_expansion): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): ReLU6()
        )
        (_conv): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=16, bias=False)
            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): ReLU6()
        )
        (_sqse): QuantizableSqueezeAndExcite(
          (_avg_pool): AdaptiveAvgPool2d(output_size=1)
          (_fc): Sequential(
            (0): LinearReLU(
              (0): Linear(in_features=16, out_features=4, bias=True)
              (1): ReLU()
            )
            (1): Identity()
            (2): Linear(in_features=4, out_features=16, bias=True)
            (3): Sequential(
              (0): DeQuantStub()
              (1): Hardsigmoid()
              (2): QuantStub()
            )
          )
          (_mul): FloatFunctional(
            (activation_post_process): Identity()
          )
        )
        (_reduce): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
        )
        (_skip_add): FloatFunctional(
          (activation_post_process): Identity()
        )
      )
      (2): QuantizableInvertedResidual(
        (_expansion): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(16, 72, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(72, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): ReLU6()
        )
        (_conv): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(72, 72, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=72, bias=False)
            (1): BatchNorm2d(72, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): ReLU6()
        )
        (_reduce): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(72, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
        )
        (_skip_add): FloatFunctional(
          (activation_post_process): Identity()
        )
      )
      (3): QuantizableInvertedResidual(
        (_expansion): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(24, 88, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(88, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): ReLU6()
        )
        (_conv): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(88, 88, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=88, bias=False)
            (1): BatchNorm2d(88, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): ReLU6()
        )
        (_reduce): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(88, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
        )
        (_skip_add): FloatFunctional(
          (activation_post_process): Identity()
        )
      )
      (4): QuantizableInvertedResidual(
        (_expansion): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(24, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_conv): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(96, 96, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), groups=96, bias=False)
            (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_sqse): QuantizableSqueezeAndExcite(
          (_avg_pool): AdaptiveAvgPool2d(output_size=1)
          (_fc): Sequential(
            (0): LinearReLU(
              (0): Linear(in_features=96, out_features=24, bias=True)
              (1): ReLU()
            )
            (1): Identity()
            (2): Linear(in_features=24, out_features=96, bias=True)
            (3): Sequential(
              (0): DeQuantStub()
              (1): Hardsigmoid()
              (2): QuantStub()
            )
          )
          (_mul): FloatFunctional(
            (activation_post_process): Identity()
          )
        )
        (_reduce): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(96, 40, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
        )
        (_skip_add): FloatFunctional(
          (activation_post_process): Identity()
        )
      )
      (5): QuantizableInvertedResidual(
        (_expansion): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_conv): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(240, 240, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=240, bias=False)
            (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_sqse): QuantizableSqueezeAndExcite(
          (_avg_pool): AdaptiveAvgPool2d(output_size=1)
          (_fc): Sequential(
            (0): LinearReLU(
              (0): Linear(in_features=240, out_features=60, bias=True)
              (1): ReLU()
            )
            (1): Identity()
            (2): Linear(in_features=60, out_features=240, bias=True)
            (3): Sequential(
              (0): DeQuantStub()
              (1): Hardsigmoid()
              (2): QuantStub()
            )
          )
          (_mul): FloatFunctional(
            (activation_post_process): Identity()
          )
        )
        (_reduce): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(240, 40, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
        )
        (_skip_add): FloatFunctional(
          (activation_post_process): Identity()
        )
      )
      (6): QuantizableInvertedResidual(
        (_expansion): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_conv): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(240, 240, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=240, bias=False)
            (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_sqse): QuantizableSqueezeAndExcite(
          (_avg_pool): AdaptiveAvgPool2d(output_size=1)
          (_fc): Sequential(
            (0): LinearReLU(
              (0): Linear(in_features=240, out_features=60, bias=True)
              (1): ReLU()
            )
            (1): Identity()
            (2): Linear(in_features=60, out_features=240, bias=True)
            (3): Sequential(
              (0): DeQuantStub()
              (1): Hardsigmoid()
              (2): QuantStub()
            )
          )
          (_mul): FloatFunctional(
            (activation_post_process): Identity()
          )
        )
        (_reduce): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(240, 40, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
        )
        (_skip_add): FloatFunctional(
          (activation_post_process): Identity()
        )
      )
      (7): QuantizableInvertedResidual(
        (_expansion): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(40, 120, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(120, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_conv): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(120, 120, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=120, bias=False)
            (1): BatchNorm2d(120, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
          (2): Hardswish()
        )
        (_sqse): QuantizableSqueezeAndExcite(
          (_avg_pool): AdaptiveAvgPool2d(output_size=1)
          (_fc): Sequential(
            (0): LinearReLU(
              (0): Linear(in_features=120, out_features=30, bias=True)
              (1): ReLU()
            )
            (1): Identity()
            (2): Linear(in_features=30, out_features=120, bias=True)
            (3): Sequential(
              (0): DeQuantStub()
              (1): Hardsigmoid()
              (2): QuantStub()
            )
          )
          (_mul): FloatFunctional(
            (activation_post_process): Identity()
          )
        )
        (_reduce): Sequential(
          (0): ConvBn2d(
            (0): Conv2d(120, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (1): Identity()
        )
        (_skip_add): FloatFunctional(
          (activation_post_process): Identity()
        )
      )
    )
  )
  (_decoder): QuantizableLRASPP(
    (_aspp_conv1): Sequential(
      (0): ConvBnReLU2d(
        (0): Conv2d(48, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU()
      )
      (1): Identity()
      (2): Identity()
    )
    (_upper2): UpsamplingBilinear2d(size=(49, 78), mode=bilinear)
    (_aspp_conv2): Sequential(
      (0): AvgPool2d(kernel_size=11, stride=(4, 4), padding=0)
      (1): Conv2d(48, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (2): Sequential(
        (0): DeQuantStub()
        (1): Sigmoid()
        (2): QuantStub()
      )
    )
    (_aspp_conv12): Conv2d(128, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (_upper12): UpsamplingBilinear2d(size=(98, 155), mode=bilinear)
    (_aspp_conv3): Conv2d(24, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (_upper123): UpsamplingBilinear2d(size=(780, 1240), mode=bilinear)
    (_mul): FloatFunctional(
      (activation_post_process): Identity()
    )
    (_add): FloatFunctional(
      (activation_post_process): Identity()
    )
  )
  (_dequant): DeQuantStub()
)
</code></pre>
<p>After model I am also performing nn.Sigmoid activation for binary pixels classification.</p>
<p>If you can help me with piece of advice don’t hesitate to reply.</p>
<p>Thanks in advance!</p>

supriyar: <p><a class="mention" href="/u/smivv">@smivv</a>, could you share the code you used to enable QAT on the model?</p>
<ol>
<li>We currently do support quantization of Sigmoid, Hardsigmoid and ReLU6. cc <a class="mention" href="/u/jerryzh168">@jerryzh168</a> to confirm</li>
<li>QAT of UpsamplingBilinear2d isn’t supported so you will have to wrap it with Quant-Dequant block.</li>
<li>Maybe looking at the code will help, we also have a tutorial here <a href="https://github.com/pytorch/vision/blob/master/references/classification/train_quantization.py" rel="noopener nofollow ugc">https://github.com/pytorch/vision/blob/master/references/classification/train_quantization.py</a> for reference</li>
<li>QAT training takes longer, due to the insertion of observers and fake quant modules in the model</li>
<li>It is possible to do QAT on GPU, but you will need to move the model to CPU before running convert.</li>
</ol>

smivv: <p>Hello <a class="mention" href="/u/supriyar">@supriyar</a>,</p>
<p>Thanks for your comment!</p>
<ol>
<li>Are you sure about Sigmoid and Hardsigmoid? Because when I look into models’ print after <code>prepate_qat</code> and I see that the only activations without <code>FakeQuantizer</code> are <code>Sigmoid</code> and <code>Hardsigmoid</code>. While <code>HardSwish</code> has <code>FakeQuantizer</code> attached:</li>
</ol>
<pre><code class="lang-auto">  (2): Hardswish(
    (activation_post_process): FakeQuantize(
      fake_quant_enabled=tensor([1], device='cuda:0', dtype=torch.uint8), observer_enabled=tensor([1], device='cuda:0', dtype=torch.uint8),            quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1,         scale=tensor([1.], device='cuda:0'), zero_point=tensor([0], device='cuda:0')
      (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)
    )
  )
........................................................
  (2): Sequential(
    (0): DeQuantStub()
    (1): Sigmoid()
    (2): QuantStub(
      (activation_post_process): FakeQuantize(
        fake_quant_enabled=tensor([1], device='cuda:0', dtype=torch.uint8), observer_enabled=tensor([1], device='cuda:0', dtype=torch.uint8),            quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1,         scale=tensor([1.], device='cuda:0'), zero_point=tensor([0], device='cuda:0')
        (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)
      )
    )
  )
</code></pre>
<ol start="2">
<li>Just did it and it didn’t help.</li>
<li>Already seen that, thanks, but nothing new in terms of code for me <img src="https://discuss.pytorch.org/images/emoji/apple/sweat.png?v=9" title=":sweat:" class="emoji" alt=":sweat:">
</li>
<li>Does it take longer in terms of epochs? Because timing for sure is increased.</li>
<li>Ok, that’s how I do that.</li>
</ol>
<p>I am using Catalyst framework, so what I am doing is that I am preparing model for QAT once stage is started. Here is my code:</p>
<pre><code class="lang-auto">class QuantizationAwareTrainingCallback(Callback):
    def __init__(
        self,
        backend: str = "fbgemm",
    ):
        super().__init__(order=CallbackOrder.Internal + 1)

        assert backend in ["fbgemm", "qnnpack"], "Unknown backend type"

        self.backend = backend

    def on_stage_start(self, runner: "IRunner") -&gt; None:
        # model is already fused
        runner.model.train()

        torch.backends.quantized.engine = self.backend
        runner.model.qconfig = torch.quantization.get_default_qat_qconfig(self.backend)
        runner.model = torch.quantization.prepare_qat(runner.model, inplace=False)

        runner.model.apply(torch.quantization.enable_observer)
        runner.model.apply(torch.quantization.enable_fake_quant)
</code></pre>

smivv: <p>Finally I made it working.</p>
<p>The real reason why it wasn’t working is because I had PyTorch 1.7 installed which is probably has not all quantized operations implemented or maybe has some bugs.</p>
<p>Now I installed PyTorch from sources using master branch and I have it training <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Thanks!</p>