Repository: quic/aimet (language: python)

Issue:
Input order flipping issue in connected graph



Discussion / hints:

Fix patch:
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/meta/connectedgraph.py b/TrainingExtensions/torch/src/python/aimet_torch/meta/connectedgraph.py
index 6457398405..8151824a42 100644
--- a/TrainingExtensions/torch/src/python/aimet_torch/meta/connectedgraph.py
+++ b/TrainingExtensions/torch/src/python/aimet_torch/meta/connectedgraph.py
@@ -83,7 +83,8 @@
 # When traced, the leaf level module for these ops may have ordering of inputs flipped. We trace into these modules in
 # order to determine the true input ordering.
 MULTI_INPUT_OPS_TO_PARSE = [aimet_modules.Add, aimet_modules.Multiply, aimet_modules.Subtract,
-                            aimet_modules.Divide, aimet_modules.Pow, aimet_modules.IndexSelect]
+                            aimet_modules.Divide, aimet_modules.Pow, aimet_modules.IndexSelect, aimet_modules.MatMul,
+                            aimet_modules.Concat, aimet_modules.Minimum, aimet_modules.DynamicConv2d]
 
 # We want to consider following operations as leaf nodes while creating op for connected graph.
 SKIP_LIST_FOR_SUBGRAPH_TRACE = [aimet_modules.StridedSlice, aimet_modules.GatherNd, aimet_modules.ScatterND,


Test patch:
diff --git a/TrainingExtensions/torch/test/python/models/test_models.py b/TrainingExtensions/torch/test/python/models/test_models.py
index 5f058f38f0..f336d1663c 100644
--- a/TrainingExtensions/torch/test/python/models/test_models.py
+++ b/TrainingExtensions/torch/test/python/models/test_models.py
@@ -1446,3 +1446,30 @@ def forward(self, inp):
         x = self.in_linear1(inp)
         x = self.linear_modlist[0](x)
         return x
+
+
+class ModelWithMultiInputOps(torch.nn.Module):
+    def __init__(self):
+        super(ModelWithMultiInputOps, self).__init__()
+        self.conv = nn.Conv2d(3, 8, kernel_size=2, stride=2, padding=0, bias=False)
+        self.bn = nn.BatchNorm2d(8)
+        self.relu = nn.ReLU(inplace=True)
+        self.register_buffer('gemm_weight', torch.ones([1, 8, 12, 24], dtype=torch.float32))
+        self.matmul = aimet_modules.MatMul()
+        self.min = aimet_modules.Minimum()
+        self.register_buffer('constant', torch.ones([1, 8, 12, 12], dtype=torch.float32))
+        self.concat = aimet_modules.Concat()
+        self.register_buffer('conv_bias', torch.zeros([3], dtype=torch.float32))
+        self.dynamic_conv = aimet_modules.DynamicConv2d(stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1)
+        self.fc = nn.Linear(24, 3)
+
+    def forward(self, inp, weight):
+        x1_conv = self.conv(inp)
+        x1_bn = self.bn(x1_conv)
+        x1 = self.relu(x1_bn)
+        x2 = self.matmul(x1, self.gemm_weight)
+        x3 = self.fc(x2)
+        x4 = self.min(self.constant, x1)
+        x5 = self.concat(x1_bn, self.constant, x1)
+        x6 = self.dynamic_conv(x5, weight, self.conv_bias)
+        return x3, x6+2, x4
diff --git a/TrainingExtensions/torch/test/python/test_quantsim_config.py b/TrainingExtensions/torch/test/python/test_quantsim_config.py
index 4a0f096b86..b0638c726c 100644
--- a/TrainingExtensions/torch/test/python/test_quantsim_config.py
+++ b/TrainingExtensions/torch/test/python/test_quantsim_config.py
@@ -2293,3 +2293,26 @@ def test_load_encodings_to_allow_modifying_quantizer_type(self):
                         assert isinstance(sim.model.conv1.param_quantizers['weight'], StaticGridPerChannelQuantizer)
                     else:
                         assert isinstance(sim.model.conv1.param_quantizers['weight'], StaticGridPerTensorQuantizer)
+
+    def test_input_order_exception_for_multi_input_ops_in_connected_graph(self):
+        model = test_models.ModelWithMultiInputOps()
+        model.eval()
+        dummy_input = (torch.randn(1, 3, 24, 24), torch.randn(3, 8, 2, 2))
+
+        sim = QuantizationSimModel(model, dummy_input)
+        # input order for the following ops must be preserved within
+        # connected graph for the quantizers to be properly enabled/disabled
+        # aimet_modules.Matmul
+        assert (not sim.model.matmul.input_quantizers[0].enabled
+                and sim.model.matmul.input_quantizers[1].enabled)
+        # aimet_modules.Minimum
+        assert (sim.model.min.input_quantizers[0].enabled
+                and not sim.model.min.input_quantizers[1].enabled)
+        # aimet_modules.Concat
+        assert (not sim.model.concat.input_quantizers[0].enabled
+                and sim.model.concat.input_quantizers[1].enabled
+                and not sim.model.concat.input_quantizers[2].enabled)
+        # aimet_modules.DynamicConv2d
+        assert (not sim.model.dynamic_conv.input_quantizers[0].enabled
+                and sim.model.dynamic_conv.input_quantizers[1].enabled
+                and sim.model.dynamic_conv.input_quantizers[2].enabled)