NXP backend: Add QAT training to AOT examples

Author: StrycekSimon

examples/nxp/aot_neutron_compile.py

@@ -1,4 +1,4 @@
-# Copyright 2024-2025 NXP
+# Copyright 2024-2026 NXP
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -39,9 +39,18 @@
     to_edge_transform_and_lower,
 )
 from executorch.extension.export_util import save_pte_program
+from torch.ao.quantization import (
+    move_exported_model_to_eval,
+    move_exported_model_to_train,
+)
 from torch.export import export
+from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e
 
-from .experimental.cifar_net.cifar_net import CifarNet, test_cifarnet_model
+from .experimental.cifar_net.cifar_net import (
+    CifarNet,
+    test_cifarnet_model,
+    train_cifarnet_model,
+)
 from .models.mobilenet_v2 import MobilenetV2
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -154,7 +163,7 @@ def get_model_and_inputs_from_name(model_name: str):
         action="store_true",
         required=False,
         default=False,
-        help="Use QAT mode for quantization (does not include QAT training)",
+        help="Use QAT mode for quantization (performs two QAT training epochs)",
     )
     parser.add_argument(
         "-s",
@@ -220,15 +229,27 @@ def get_model_and_inputs_from_name(model_name: str):
 
     # 3. Quantize if required
     if args.quantize:
-        if calibration_inputs is None:
-            logging.warning(
-                "No calibration inputs available, using the example inputs instead"
-            )
-            calibration_inputs = example_inputs
-        quantizer = NeutronQuantizer(neutron_target_spec, args.use_qat)
-        module = calibrate_and_quantize(
-            module, calibration_inputs, quantizer, is_qat=args.use_qat
-        )
+        quantizer = NeutronQuantizer(neutron_target_spec, is_qat=args.use_qat)
+        if args.use_qat:
+            match args.model_name:
+                case "cifar10":
+                    print("Starting two epochs of QAT training with CifarNet model...")
+                    module = prepare_qat_pt2e(module, quantizer)
+                    module = move_exported_model_to_train(module)
+                    module = train_cifarnet_model(module, num_epochs=2)
+                    module = move_exported_model_to_eval(module)
+                    module = convert_pt2e(module)
+                case _:
+                    raise ValueError(
+                        f"QAT training is not supported for model '{args.model_name}'"
+                    )
+        else:
+            if calibration_inputs is None:
+                logging.warning(
+                    "No calibration inputs available, using the example inputs instead"
+                )
+                calibration_inputs = example_inputs
+            module = calibrate_and_quantize(module, calibration_inputs, quantizer)
 
     if args.so_library is not None:
         logging.debug(f"Loading libraries: {args.so_library}")

examples/nxp/experimental/cifar_net/cifar_net.py

@@ -1,4 +1,4 @@
-# Copyright 2024-2025 NXP
+# Copyright 2024-2026 NXP
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -145,34 +145,11 @@ def get_model(
         cifar_net.load_state_dict(torch.load(state_dict_file, weights_only=True))
 
     if train:
-        # Train the model.
-        criterion = nn.CrossEntropyLoss()
-        optimizer = optim.SGD(cifar_net.parameters(), lr=0.0005, momentum=0.6)
-        train_loader = get_train_loader(batch_size)
-
-        for epoch in range(num_epochs):
-            running_loss = 0.0
-            for i, data in enumerate(train_loader, 0):
-                # get the inputs; data is a list of [inputs, labels]
-                inputs, labels = data
-
-                # zero the parameter gradients
-                optimizer.zero_grad()
-
-                # forward + backward + optimize
-                outputs = cifar_net(inputs)
-                loss = criterion(outputs, labels)
-                loss.backward()
-                optimizer.step()
-
-                # print statistics
-                running_loss += loss.item()
-                if i % 2000 == 1999:  # print every 2000 mini-batches
-                    print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
-                    running_loss = 0.0
-
-        logger.info("Finished training.")
-        if state_dict_file is not None and train:
+        cifar_net = train_cifarnet_model(
+            cifar_net=cifar_net, batch_size=batch_size, num_epochs=num_epochs
+        )
+
+        if state_dict_file is not None:
             logger.info(f"Saving the trained weights in `{state_dict_file}`.")
             torch.save(cifar_net.state_dict(), state_dict_file)
 
@@ -189,6 +166,40 @@ def get_cifarnet_calibration_data(num_images: int = 100) -> tuple[torch.Tensor]:
     return (tensor,)
 
 
+def train_cifarnet_model(
+    cifar_net: nn.Module | torch.fx.GraphModule,
+    batch_size: int = 1,
+    num_epochs: int = 1,
+) -> nn.Module:
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(cifar_net.parameters(), lr=0.0001, momentum=0.6)
+    train_loader = get_train_loader(batch_size)
+
+    for epoch in range(num_epochs):
+        running_loss = 0.0
+        for i, data in enumerate(train_loader, 0):
+            # get the inputs; data is a list of [inputs, labels]
+            inputs, labels = data
+
+            # zero the parameter gradients
+            optimizer.zero_grad()
+
+            # forward + backward + optimize
+            outputs = cifar_net(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            # print statistics
+            running_loss += loss.item()
+            if i % 2000 == 1999:  # print every 2000 mini-batches
+                print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
+                running_loss = 0.0
+
+    logger.info("Finished training.")
+    return cifar_net
+
+
 def test_cifarnet_model(cifar_net: nn.Module, batch_size: int = 1) -> float:
     """Test the CifarNet model on the CifarNet10 testing dataset and return the accuracy.