[Fix] onnx_qnn backend's output needs a dequant. (#126)

* [Fix] onnx_qnn backend's output needs a dequant.

* [Test] add unit test for onnx_qnn

Co-authored-by: fanyunqian <[email protected]>

Author: PannenetsF

mqbench/deploy/deploy_onnx_qnn.py

@@ -210,15 +210,15 @@ def search_and_replace_input(next_node, name, new_name):
                     if prev_node != 'INPUT_TOKEN' and prev_node.op_type in self.qlinear_op_type and \
                             next_node != 'OUTPUT_TOKEN' and next_node.op_type in self.qlinear_op_type:
                         search_and_replace_input(next_node, node.output[0], node.input[0])
-                    elif prev_node != 'INPUT_TOKEN' and prev_node.op_type in self.qlinear_op_type:
+                    elif prev_node != 'INPUT_TOKEN' and prev_node.op_type in self.qlinear_op_type and \
+                            next_node == 'OUTPUT_TOKEN':
                         if dequantize_node is None:
                             output_value_info = [f'{node.output[0]}_DequantizeLinear']
                             dequantize_node = onnx.helper.make_node("DequantizeLinear",
                                                                     node.input[0:3],
                                                                     output_value_info,
                                                                     ('input' if prev_node == 'INPUT_TOKEN' else prev_node.name) + '_dequantized')
                             self.onnx_model.insert_node_purely(dequantize_node)
-                        search_and_replace_input(next_node, node.output[0], dequantize_node.output[0])
                     else:
                         if quantize_node is None:
                             output_value_info = [f'{node.output[0]}_QuantizeLinear']

test/backend/test_backend.py

@@ -5,6 +5,7 @@
 from mqbench.convert_deploy import convert_deploy
 from mqbench.utils.state import enable_calibration, enable_quantization
 
+from .test_model.unet import UNet
 from ..version import GITHUB_RES
 
 
@@ -117,7 +118,7 @@ def test_quantize_vitis(self):
         else:
             pass
 
-    def test_quantize_onnxqnn(self):
+    def test_quantize_onnxqnn_1(self):
         model_to_quantize = torch.hub.load(GITHUB_RES, 'resnet18', pretrained=False)
         dummy_input = torch.randn(2, 3, 224, 224, device='cpu')
         model_to_quantize.train()
@@ -130,6 +131,19 @@ def test_quantize_onnxqnn(self):
         model_prepared.eval()
         convert_deploy(model_prepared, BackendType.ONNX_QNN, {'x': [1, 3, 224, 224]}, model_name='resnet18_onnx_qnn.onnx')
 
+    def test_quantize_onnxqnn_2(self):
+        model_to_quantize = UNet(3, 2)
+        dummy_input = torch.randn(2, 3, 224, 224, device='cpu')
+        model_to_quantize.train()
+        model_prepared = prepare_by_platform(model_to_quantize, BackendType.ONNX_QNN)
+        enable_calibration(model_prepared)
+        model_prepared(dummy_input)
+        enable_quantization(model_prepared)
+        loss = model_prepared(dummy_input).sum()
+        loss.backward()
+        model_prepared.eval()
+        convert_deploy(model_prepared, BackendType.ONNX_QNN, {'x': [1, 3, 224, 224]}, model_name='resnet18_onnx_qnn.onnx')
+
     def test_quantize_ppl_cuda(self):
         import numpy as np
         model_to_quantize = torch.hub.load(GITHUB_RES, 'resnet18', pretrained=False)

test/backend/test_model/unet.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2022 Carl Zeiss AG – All Rights Reserved.
+# ZEISS, ZEISS.com are registered trademarks of Carl Zeiss AG
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['UNet']
+
+class UNet(nn.Module):
+    def __init__(
+        self, num_channels, num_classes, depth=4, initial_filter_count=64, bilinear=True
+    ):
+        super(UNet, self).__init__()
+
+        self.num_channels = num_channels
+        self.num_classes = num_classes
+        self.depth = depth
+        self.initial_filter_count = initial_filter_count
+        self.bilinear = bilinear
+
+        factor = 2 if bilinear else 1
+
+        filter_count = initial_filter_count
+
+        encoder_blocks = []
+        encoder_blocks.append(DoubleConv(num_channels, filter_count))
+        for d in range(depth):
+            if d < depth - 1:
+                encoder_blocks.append(Down(filter_count, 2 * filter_count))
+            else:
+                encoder_blocks.append(Down(filter_count, (2 * filter_count) // factor))
+            filter_count *= 2
+        self.encoder_blocks = nn.Sequential(*encoder_blocks)
+
+        decoder_blocks = []
+        for d in range(depth):
+            if d < depth - 1:
+                decoder_blocks.append(
+                    Up(filter_count, filter_count // 2 // factor, bilinear)
+                )
+            else:
+                decoder_blocks.append(Up(filter_count, filter_count // 2, bilinear))
+            filter_count //= 2
+        self.decoder_blocks = nn.Sequential(*decoder_blocks)
+
+        self.outc = OutputConvolution(filter_count, num_classes)
+
+    def forward(self, x):
+        xs = []
+        for encoder_block in self.encoder_blocks:
+            x = encoder_block(x)
+            xs.append(x)
+
+        xs.reverse()
+        xs = xs[1:]
+
+        for decoder_block, x_skip in zip(self.decoder_blocks, xs):
+            x = decoder_block(x, x_skip)
+
+        logits = self.outc(x)
+
+        return logits
+
+
+class DoubleConv(nn.Module):
+    """Module combining Conv -> BN -> ReLU -> Conv -> BN -> ReLU."""
+
+    def __init__(
+        self, num_input_channels, num_output_channels, num_middle_channels=None
+    ):
+        super().__init__()
+
+        if not num_middle_channels:
+            num_middle_channels = num_output_channels
+
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(
+                num_input_channels,
+                num_middle_channels,
+                kernel_size=3,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_middle_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(
+                num_middle_channels,
+                num_output_channels,
+                kernel_size=3,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_output_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """Module combining downscaling and DoubleConvolution."""
+
+    def __init__(self, num_input_channels, num_output_channels):
+        super().__init__()
+
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2), DoubleConv(num_input_channels, num_output_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Module combining upscaling and DoubleConvolution."""
+
+    def __init__(self, num_input_channels, num_output_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
+            self.conv = DoubleConv(
+                num_input_channels, num_output_channels, num_input_channels // 2
+            )
+        else:
+            self.up = nn.ConvTranspose2d(
+                num_input_channels, num_input_channels // 2, kernel_size=2, stride=2
+            )
+            self.conv = DoubleConv(num_input_channels, num_output_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        x = torch.cat([x2, x1], dim=1)
+
+        return self.conv(x)
+
+
+class OutputConvolution(nn.Module):
+    def __init__(self, num_input_channels, num_output_channels):
+        super(OutputConvolution, self).__init__()
+
+        self.conv = nn.Conv2d(num_input_channels, num_output_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)