From f25ccdadaa5418bcdeeafefa21515ddf5f5d1d06 Mon Sep 17 00:00:00 2001 From: Christophe Favergeon Date: Thu, 2 Jul 2026 13:48:11 +0200 Subject: [PATCH 1/5] cmsis pack test: first version exporting all test data in the pte Restricted to a small set of operators to test. Will be extended to all operators in the future. pte now contains atol, rtol, input and output test patterns. --- .../test/all_ops/generate_test_models.py | 93 +++++++++++++------ .../arm/cmsis_pack/test/all_ops/op_recipes.py | 11 +-- 2 files changed, 70 insertions(+), 34 deletions(-) diff --git a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py index 99e6dd2e4a7..95b09019eec 100755 --- a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py +++ b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py @@ -32,6 +32,7 @@ sys.path.insert(0, str(_SCRIPTS)) import op_recipes # type: ignore[import-not-found] # noqa: E402 + from op_guards import ( # type: ignore[import-not-found] # noqa: E402 discover_components, ) @@ -73,18 +74,49 @@ def _flatten_outputs(out) -> list: return flat raise TypeError(f"unsupported output type {type(out)}") +def _get_number_of_outputs(outputs) -> int: + if isinstance(outputs, torch.Tensor): + return 1 + elif isinstance(outputs, (tuple, list)): + return len(outputs) + else: + raise TypeError(f"unsupported output type {type(outputs)}") + -def _export_portable(model: torch.nn.Module, inputs: tuple,recipe: op_recipes.Recipe, display_quantized_values: bool = False) -> tuple[bytes, float, float]: +def _mk_metadata(inputs: tuple, outputs: tuple | torch.Tensor, atol: float, rtol: float) -> dict: + metadata = {} + metadata["nb_inputs"] = len(inputs) + metadata["nb_outputs"] = _get_number_of_outputs(outputs) + metadata["atol"] = atol + metadata["rtol"] = rtol + + for i, t in enumerate(inputs): + metadata[f"input_{i}"] = t + if isinstance(outputs, torch.Tensor): + metadata["output_0"] = outputs + else: + for i,t in enumerate(outputs): + metadata[f"output_{i}"] = t + + return metadata + +def _export_portable(model: torch.nn.Module, inputs: tuple,recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> bytes: from executorch.exir import EdgeCompileConfig, to_edge model = model.eval() exported = torch.export.export(model, inputs, strict=True) + + expected = model(*inputs) + metadata = _mk_metadata(inputs, expected,recipe.atol, recipe.rtol) + if display_metadata: + print(metadata) # The pack ships the full portable op set, including ops outside the Core # ATen opset (bitwise shifts, unfold, var_mean.correction, ...), so skip the # core-ATen IR validity gate; to_executorch still lowers to the .out kernels. - edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False)) + edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False), + constant_methods=metadata) program = edge.to_executorch() - return bytes(program.buffer), recipe.atol, recipe.rtol + return bytes(program.buffer) def _compute_test_threshold(actual, expected): abs_err = (actual - expected).abs() @@ -99,8 +131,8 @@ def _compute_test_threshold(actual, expected): rtol = 0.0 return atol, rtol -def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.Recipe, display_quantized_values: bool = False) -> tuple[bytes, float, float]: - from executorch.backends.cortex_m.passes.cortex_m_pass_manager import ( +def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> bytes: + from debug_cortex_m.passes.cortex_m_pass_manager import ( CortexMPassManager, ) from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer @@ -111,6 +143,7 @@ def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.R model = model.eval() expected = model(*inputs) + captured = torch.export.export(model, inputs, strict=True).module() prepared = prepare_pt2e(captured, CortexMQuantizer()) prepared(*inputs) # calibrate @@ -119,10 +152,16 @@ def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.R actual = quantized(*inputs) if display_quantized_values: print("=== quantized values ===") - print("Expected:", expected) - print("Actual:", actual) + print("Expected:", expected,expected.shape) + print("Actual:", actual,actual.shape) atol,rtol = _compute_test_threshold(actual, expected) + + metadata = _mk_metadata(inputs, expected,atol, rtol) + if display_metadata: + print(metadata) + exported = torch.export.export(quantized, inputs, strict=True) + edge = to_edge_transform_and_lower( exported, compile_config=EdgeCompileConfig( preserve_ops=[ @@ -134,12 +173,13 @@ def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.R ], _check_ir_validity=False, _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default], - ) + ), + constant_methods=metadata ) edge._edge_programs["forward"] = CortexMPassManager( edge.exported_program() ).transform() - return bytes(edge.to_executorch().buffer),atol,rtol + return bytes(edge.to_executorch().buffer) def main() -> None: @@ -165,6 +205,11 @@ def main() -> None: action="store_true", help="display quantized values during export", ) + parser.add_argument( + "--display-metadata", + action="store_true", + help="display metadata during export", + ) args = parser.parse_args() source_dir = Path(args.source_dir) @@ -193,36 +238,30 @@ def main() -> None: for component in components: key = (component.category, component.name) + if component.name != "quantized_max_pool2d" and component.name != "quantized_conv2d" and component.name != "add": + skipped.append((component.category, component.name, "For debug")) + continue + if key in op_recipes.SKIPS: skipped.append((component.category, component.name, op_recipes.SKIPS[key])) continue recipe = op_recipes.RECIPES[key] cat_id = component.category.lower().replace("-", "_") - op_dir = out_dir / f"{cat_id}__{component.name}" + op_name = out_dir / f"{cat_id}__{component.name}" + op_pte = op_name.with_suffix(".pte") try: - if args.display_quantized_values: + if args.display_quantized_values or args.display_metadata: print(f"=== exporting {component.category}/{component.name} ===") model, inputs = recipe.make() - reference = model.eval()(*inputs) - pte,atol,rtol = exporters[component.category](model, inputs, recipe,display_quantized_values=args.display_quantized_values) - op_dir.mkdir(parents=True, exist_ok=True) - (op_dir / "model.pte").write_bytes(pte) - in_specs = [ - _save_tensor(t, op_dir / f"input_{i}.bin") for i, t in enumerate(inputs) - ] - out_specs = [ - _save_tensor(t, op_dir / f"expected_{i}.bin") - for i, t in enumerate(_flatten_outputs(reference)) - ] + + pte = exporters[component.category](model, inputs, recipe,display_quantized_values=args.display_quantized_values,display_metadata=args.display_metadata) + op_pte.write_bytes(pte) + manifest.append( { "op": component.name, "category": component.category, - "dir": op_dir.name, - "atol": atol, - "rtol": rtol, - "inputs": [s.__dict__ for s in in_specs], - "outputs": [s.__dict__ for s in out_specs], + "name": op_name.name } ) except Exception as exc: # noqa: BLE001 - report, don't abort the sweep diff --git a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py index 0b3abe1ecfd..4ec10d01567 100644 --- a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py +++ b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py @@ -55,14 +55,11 @@ def forward(self, *args): return self.fn(*args) -def _ramp(lo: float, hi: float, shape: tuple, channel_last=False) -> torch.Tensor: +def _ramp(lo: float, hi: float, shape: tuple) -> torch.Tensor: n = 1 for s in shape: n *= s - if channel_last: - return torch.linspace(lo, hi, n).reshape(shape).contiguous(memory_format=torch.channels_last) - else: - return torch.linspace(lo, hi, n).reshape(shape) + return torch.linspace(lo, hi, n).reshape(shape) def _reg(category: str, name: str, make: Callable, **kw) -> None: @@ -578,11 +575,11 @@ def _cm(name: str, make: Callable, **kw): ) _cm( "quantized_avg_pool2d", - lambda: (torch.nn.AvgPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8), channel_last=True),)), + lambda: (torch.nn.AvgPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8)))), ) _cm( "quantized_max_pool2d", - lambda: (torch.nn.MaxPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8), channel_last=True),)), + lambda: (torch.nn.MaxPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8)),)), ) _cm( "quantized_batch_matmul", From 5296f696048ba83d0f7b8888c3af2da636c53b77 Mon Sep 17 00:00:00 2001 From: Christophe Favergeon Date: Fri, 3 Jul 2026 08:37:36 +0200 Subject: [PATCH 2/5] CMSIS pack tests : Updated recipes to select right kernels Only focusing on quantized_pool_xxx and quantized_conv2d for this commit. --- .../test/all_ops/generate_test_models.py | 48 +++++++------------ .../arm/cmsis_pack/test/all_ops/op_recipes.py | 16 +++++-- 2 files changed, 30 insertions(+), 34 deletions(-) diff --git a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py index 95b09019eec..90913fcaba5 100755 --- a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py +++ b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py @@ -45,34 +45,6 @@ class TensorSpec: shape: list -def _save_tensor(t: torch.Tensor, path: Path) -> TensorSpec: - src = t.detach().cpu() - # The firmware copies these bytes directly into ExecuTorch tensor storage. - # For channels_last tensors, that storage is NHWC even though the logical - # tensor shape remains NCHW. - if ( - src.dim() == 4 - and src.is_contiguous(memory_format=torch.channels_last) - and not src.is_contiguous() - ): - storage = src.permute(0, 2, 3, 1).contiguous() - else: - storage = src.contiguous() - path.write_bytes(storage.numpy().tobytes()) - return TensorSpec( - file=path.name, dtype=str(src.dtype).replace("torch.", ""), shape=list(src.shape) - ) - - -def _flatten_outputs(out) -> list: - if isinstance(out, torch.Tensor): - return [out] - if isinstance(out, (tuple, list)): - flat = [] - for o in out: - flat.extend(_flatten_outputs(o)) - return flat - raise TypeError(f"unsupported output type {type(out)}") def _get_number_of_outputs(outputs) -> int: if isinstance(outputs, torch.Tensor): @@ -90,13 +62,27 @@ def _mk_metadata(inputs: tuple, outputs: tuple | torch.Tensor, atol: float, rtol metadata["atol"] = atol metadata["rtol"] = rtol + channel_last = False + # We assume that when one input tensor is channel_last, all others too + # This assumption is true for the operators tested for i, t in enumerate(inputs): metadata[f"input_{i}"] = t + if t.is_contiguous(memory_format=torch.channels_last): + channel_last = True if isinstance(outputs, torch.Tensor): metadata["output_0"] = outputs else: for i,t in enumerate(outputs): metadata[f"output_{i}"] = t + + # Input / outputs are exported as channel_first. + # So we need to store the memory format information in the metadata so that + # we can change the input / output tensors to channel_last in the C++ + # tests. + # It is a workaround for what looks like a bug in the to_edge_transform_and_lower + # function, which does not preserve the memory format of the input / output tensors + # exported as constant_methods. + metadata["channel_last"] = channel_last return metadata @@ -238,9 +224,9 @@ def main() -> None: for component in components: key = (component.category, component.name) - if component.name != "quantized_max_pool2d" and component.name != "quantized_conv2d" and component.name != "add": - skipped.append((component.category, component.name, "For debug")) - continue + #if component.name != "quantized_max_pool2d" and component.name != "quantized_conv2d" and component.name != "quantized_avg_pool2d": + # skipped.append((component.category, component.name, "For debug")) + # continue if key in op_recipes.SKIPS: skipped.append((component.category, component.name, op_recipes.SKIPS[key])) diff --git a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py index 4ec10d01567..e83d175e628 100644 --- a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py +++ b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py @@ -522,6 +522,16 @@ def _red(name: str, fn: Callable, **kw): ) +# Use instead of torch.nn.Conv2d to set the weights +class CortexMConv2D(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.conv = torch.nn.Conv2d(*args, **kwargs, bias=False) + self.conv.weight.data.fill_(1.0) + + def forward(self, x): + return self.conv(x) + # -------------------------------------------------------------------------- # Cortex-M ops (quantized; exercised via the CortexM export flow). Inputs are # float; the CortexMQuantizer + passes lower them to cortex_m::* kernels. @@ -563,7 +573,7 @@ def _cm(name: str, make: Callable, **kw): ) _cm( "quantized_conv2d", - lambda: (torch.nn.Conv2d(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),)), + lambda: (CortexMConv2D(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)), ) _cm( "quantized_depthwise_conv2d", @@ -575,11 +585,11 @@ def _cm(name: str, make: Callable, **kw): ) _cm( "quantized_avg_pool2d", - lambda: (torch.nn.AvgPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8)))), + lambda: (torch.nn.AvgPool2d(kernel_size=2, stride=2), (_ramp(0, 15, (1, 1, 4, 4)),)), ) _cm( "quantized_max_pool2d", - lambda: (torch.nn.MaxPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8)),)), + lambda: (torch.nn.MaxPool2d(kernel_size=2, stride=2), (_ramp(-50, 50, (1, 1, 6, 6)),)), ) _cm( "quantized_batch_matmul", From 9ae4c362c15a1a00a9184d5ae8559f8b80088d69 Mon Sep 17 00:00:00 2001 From: Christophe Favergeon Date: Fri, 3 Jul 2026 09:17:44 +0200 Subject: [PATCH 3/5] CMSIS pack tests: Enable new quantized kernels and allow to dump graph before export --- .../test/all_ops/generate_test_models.py | 24 ++++++++++++++----- .../arm/cmsis_pack/test/all_ops/op_recipes.py | 4 ++-- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py index 90913fcaba5..c61bb760f4c 100755 --- a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py +++ b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py @@ -86,7 +86,7 @@ def _mk_metadata(inputs: tuple, outputs: tuple | torch.Tensor, atol: float, rtol return metadata -def _export_portable(model: torch.nn.Module, inputs: tuple,recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> bytes: +def _export_portable(model: torch.nn.Module, inputs: tuple,recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> ExecutorchProgramManager: from executorch.exir import EdgeCompileConfig, to_edge model = model.eval() @@ -102,7 +102,7 @@ def _export_portable(model: torch.nn.Module, inputs: tuple,recipe: op_recipes.Re edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False), constant_methods=metadata) program = edge.to_executorch() - return bytes(program.buffer) + return program def _compute_test_threshold(actual, expected): abs_err = (actual - expected).abs() @@ -117,7 +117,7 @@ def _compute_test_threshold(actual, expected): rtol = 0.0 return atol, rtol -def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> bytes: +def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> ExecutorchProgramManager: from debug_cortex_m.passes.cortex_m_pass_manager import ( CortexMPassManager, ) @@ -165,7 +165,7 @@ def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.R edge._edge_programs["forward"] = CortexMPassManager( edge.exported_program() ).transform() - return bytes(edge.to_executorch().buffer) + return edge.to_executorch() def main() -> None: @@ -196,6 +196,14 @@ def main() -> None: action="store_true", help="display metadata during export", ) + # Model explorer (with pte extension) or Netron are not able to display + # all the .pte in the right way. + # So, another way to check is to print the final graph. + parser.add_argument( + "--display-graph", + action="store_true", + help="display graph during export", + ) args = parser.parse_args() source_dir = Path(args.source_dir) @@ -224,7 +232,8 @@ def main() -> None: for component in components: key = (component.category, component.name) - #if component.name != "quantized_max_pool2d" and component.name != "quantized_conv2d" and component.name != "quantized_avg_pool2d": + # For debugging. To select only the nodes under investigation + #if component.name != "convolution" and component.name != "embedding": # skipped.append((component.category, component.name, "For debug")) # continue @@ -241,7 +250,10 @@ def main() -> None: model, inputs = recipe.make() pte = exporters[component.category](model, inputs, recipe,display_quantized_values=args.display_quantized_values,display_metadata=args.display_metadata) - op_pte.write_bytes(pte) + if args.display_graph: + _ = pte.exported_program("forward").graph_module.print_readable() + + op_pte.write_bytes(bytes(pte.buffer)) manifest.append( { diff --git a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py index e83d175e628..310d5bb35e5 100644 --- a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py +++ b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py @@ -577,11 +577,11 @@ def _cm(name: str, make: Callable, **kw): ) _cm( "quantized_depthwise_conv2d", - lambda: (torch.nn.Conv2d(4, 4, 3, groups=4), (_ramp(-1, 1, (1, 4, 8, 8)),)), + lambda: (torch.nn.Conv2d(4, 4, 3, groups=4), (_ramp(1, 5, (1, 4, 8, 8)).to(memory_format=torch.channels_last),)), ) _cm( "quantized_transpose_conv2d", - lambda: (torch.nn.ConvTranspose2d(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),)), + lambda: (torch.nn.ConvTranspose2d(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)), ) _cm( "quantized_avg_pool2d", From c7384b3d8208f2f3ab70572363bc7c565579eb13 Mon Sep 17 00:00:00 2001 From: Christophe Favergeon Date: Fri, 3 Jul 2026 09:30:37 +0200 Subject: [PATCH 4/5] Remove reference to a debug python package --- backends/arm/cmsis_pack/test/all_ops/generate_test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py index c61bb760f4c..0213505949e 100755 --- a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py +++ b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py @@ -118,7 +118,7 @@ def _compute_test_threshold(actual, expected): return atol, rtol def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> ExecutorchProgramManager: - from debug_cortex_m.passes.cortex_m_pass_manager import ( + from executorch.backends.cortex_m.passes.cortex_m_pass_manager import ( CortexMPassManager, ) from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer From 0ac9180c05fc07e1844bb3fb05d615e681bbe44e Mon Sep 17 00:00:00 2001 From: Christophe Favergeon Date: Mon, 6 Jul 2026 14:57:35 +0200 Subject: [PATCH 5/5] cmsis_pack tests : remove random weights in operators For the tests to be reproducible, we need to remove the random weight initializations in some operators. This is done by setting the weights (and bias) to a constant value (1.0) in the operator recipes. --- .../arm/cmsis_pack/test/all_ops/op_recipes.py | 75 +++++++++++++++---- 1 file changed, 60 insertions(+), 15 deletions(-) diff --git a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py index 310d5bb35e5..185ae5aae46 100644 --- a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py +++ b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py @@ -43,6 +43,60 @@ class Recipe: # Explicit skips keyed the same way: ops we deliberately do not execute. SKIPS: dict[tuple[str, str], str] = {} +# torch.nn.Conv2d weights are initialized randomly +# To have reproducible results, we create a custom Conv2d module +# that initializes the weights to 1.0 +class FixedParametersConv2D(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.conv = torch.nn.Conv2d(*args, **kwargs, bias=False) + self.conv.weight.data.fill_(1.0) + + def forward(self, x): + return self.conv(x) + + +# torch.nn.Conv2d weights are initialized randomly +# To have reproducible results, we create a custom Conv2d module +# that initializes the weights to 1.0 +class FixedParametersEmbedding(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.embedding = torch.nn.Embedding(*args, **kwargs) + self.embedding.weight.data.fill_(1.0) + + def forward(self, x): + return self.embedding(x) + +# torch.nn.Linear weights are initialized randomly +# To have reproducible results, we create a custom Linear module +# that initializes the weights to 1.0 +class FixedParametersLinear(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.linear = torch.nn.Linear(*args, **kwargs) + self.linear.weight.data.fill_(1.0) + if self.linear.bias is not None: + self.linear.bias.data.fill_(1.0) + + def forward(self, x): + return self.linear(x) + +# torch.nn.ConvTranspose2d weights are initialized randomly +# To have reproducible results, we create a custom ConvTranspose2d module +# that initializes the weights to 1.0 +class FixedParametersConvTranspose2d(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.conv_transpose = torch.nn.ConvTranspose2d(*args, **kwargs) + self.conv_transpose.weight.data.fill_(1.0) + if self.conv_transpose.bias is not None: + self.conv_transpose.bias.data.fill_(1.0) + + def forward(self, x): + return self.conv_transpose(x) + + class _Fn(torch.nn.Module): """Wrap a plain callable as a module so torch.export can trace it.""" @@ -286,7 +340,7 @@ def _red(name: str, fn: Callable, **kw): # Conv / pool / norm # -------------------------------------------------------------------------- _portable( - "convolution", lambda: (torch.nn.Conv2d(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),)) + "convolution", lambda: (FixedParametersConv2D(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),)) ) _portable("avg_pool2d", lambda: (torch.nn.AvgPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8)),))) _portable( @@ -418,7 +472,7 @@ def _red(name: str, fn: Callable, **kw): ) _portable( "embedding", - lambda: (torch.nn.Embedding(5, 3), (torch.tensor([[0, 2], [1, 4]]),)), + lambda: (FixedParametersEmbedding(5, 3), (torch.tensor([[0, 2], [1, 4]]),)), ) _portable( "scatter_add", @@ -522,15 +576,6 @@ def _red(name: str, fn: Callable, **kw): ) -# Use instead of torch.nn.Conv2d to set the weights -class CortexMConv2D(torch.nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - self.conv = torch.nn.Conv2d(*args, **kwargs, bias=False) - self.conv.weight.data.fill_(1.0) - - def forward(self, x): - return self.conv(x) # -------------------------------------------------------------------------- # Cortex-M ops (quantized; exercised via the CortexM export flow). Inputs are @@ -558,7 +603,7 @@ def _cm(name: str, make: Callable, **kw): ) _cm( "quantized_linear", - lambda: (torch.nn.Linear(8, 4, bias=False), (_ramp(-2, 2, (1, 8)),)), + lambda: (FixedParametersLinear(8, 4, bias=False), (_ramp(-2, 2, (1, 8)),)), ) _cm( "softmax", lambda: (_Fn(lambda x: torch.softmax(x, dim=1)), (_ramp(-2, 2, (1, 8)),)) @@ -573,15 +618,15 @@ def _cm(name: str, make: Callable, **kw): ) _cm( "quantized_conv2d", - lambda: (CortexMConv2D(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)), + lambda: (FixedParametersConv2D(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)), ) _cm( "quantized_depthwise_conv2d", - lambda: (torch.nn.Conv2d(4, 4, 3, groups=4), (_ramp(1, 5, (1, 4, 8, 8)).to(memory_format=torch.channels_last),)), + lambda: (FixedParametersConv2D(4, 4, 3, groups=4), (_ramp(1, 5, (1, 4, 8, 8)).to(memory_format=torch.channels_last),)), ) _cm( "quantized_transpose_conv2d", - lambda: (torch.nn.ConvTranspose2d(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)), + lambda: (FixedParametersConvTranspose2d(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)), ) _cm( "quantized_avg_pool2d",