NXP backend: Disable training mode and replace deprecated call (#13756)

### Summary
Ensure the input models are always in evaluation mode.
Additionally, a deprecated export call was replaced by an up-to-date
one.

### Test plan
Correct function is tested by most of the existing tests.

Author: MartinPavella

backends/nxp/tests/executorch_pipeline.py

@@ -96,9 +96,10 @@ def to_quantized_edge_program(
 
     example_input = calibration_inputs[0]
 
-    exir_program_aten = torch.export.export_for_training(
-        model, example_input, strict=True
-    )
+    # Make sure the model is in the evaluation mode.
+    model.eval()
+
+    exir_program_aten = torch.export.export(model, example_input, strict=True)
 
     exir_program_aten__module_quant = _quantize_model(
         exir_program_aten.module(), calibration_inputs
@@ -147,5 +148,9 @@ def to_edge_program(
     calibration_inputs = get_random_calibration_inputs(to_model_input_spec(input_spec))
 
     example_input = calibration_inputs[0]
+
+    # Make sure the model is in the evaluation mode.
+    model.eval()
+
     exir_program = torch.export.export(model, example_input)
     return exir.to_edge(exir_program)

backends/nxp/tests/test_batch_norm_fusion.py

@@ -95,7 +95,7 @@ def test_batch_norm_conv_fusing(bias: bool, input_shape: list[int]):
     example_input = (torch.ones(*input_shape),)
 
     module = ConvBatchNormModule(bias, len(input_shape), 4)
-    program = torch.export.export_for_training(module, example_input, strict=True)
+    program = torch.export.export(module, example_input, strict=True)
     og_module = program.module()
 
     pm = NeutronAtenPassManager()
@@ -129,7 +129,7 @@ def test_batch_norm_linear_fusing(bias: bool):
     example_input = (torch.ones(*input_shape),)
 
     module = LinearBatchNormModule(bias, 4, input_shape[-1], input_shape[1])
-    program = torch.export.export_for_training(module, example_input, strict=True)
+    program = torch.export.export(module, example_input, strict=True)
     og_module = program.module()
 
     pm = NeutronAtenPassManager()

backends/nxp/tests/test_quantizer.py

@@ -23,9 +23,7 @@ def test_quantizer_conv2d():
 
     example_input = (torch.ones(1, 4, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -64,9 +62,7 @@ def test_quantizer_linear():
 
     example_input = (torch.ones(10, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -105,9 +101,7 @@ def test_quantizer_maxpool2d():
 
     example_input = (torch.ones(1, 8, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -143,9 +137,7 @@ def test_quantizer_softmax():
 
     example_input = (torch.ones(1, 10),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -182,9 +174,7 @@ def test_quantizer_single_maxpool2d():
 
     example_input = (torch.ones(1, 4, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -206,9 +196,7 @@ def test_quantizer_conv2d_relu():
 
     example_input = (torch.ones(1, 4, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -231,9 +219,7 @@ def test_quantizer_conv2d_avg_pool2d():
 
     example_input = (torch.ones(1, 4, 16, 16),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -256,9 +242,7 @@ def test_quantizer_conv2d_permute():
 
     example_input = (torch.ones(1, 4, 16, 16),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -285,9 +269,7 @@ def test_multiple_shared_spec_ops_in_row():
 
     example_input = (torch.ones(1, 3, 64, 64),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -321,9 +303,7 @@ def test_quantizers_order_invariance():
     example_input = (torch.ones(1, 4, 64, 64),)
     quantizer = NeutronQuantizer()
 
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     m = prepare_pt2e(deepcopy(graph_module), quantizer)
     m(*example_input)

examples/nxp/aot_neutron_compile.py

@@ -15,7 +15,6 @@
 import executorch.kernels.quantized  # noqa F401
 
 import torch
-
 from executorch.backends.nxp.backend.ir.edge_passes.remove_io_quant_ops_pass import (
     RemoveIOQuantOpsPass,
 )
@@ -24,14 +23,12 @@
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 from executorch.examples.models.model_factory import EagerModelFactory
-
 from executorch.exir import (
     EdgeCompileConfig,
     ExecutorchBackendConfig,
     to_edge_transform_and_lower,
 )
 from executorch.extension.export_util import save_pte_program
-
 from torch.export import export
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
@@ -227,9 +224,7 @@ def _get_batch_size(data):
     model = model.eval()
 
     # 2. Export the model to ATEN
-    exported_program = torch.export.export_for_training(
-        model, example_inputs, strict=True
-    )
+    exported_program = torch.export.export(model, example_inputs, strict=True)
 
     module = exported_program.module()